Vulkan texture direct mapping

Better handling for row and slice pitch when writing/reading to/from Vulkan images

Source/BansheeCore/Include/BsPixelData.h

@@ -169,28 +169,20 @@ namespace bs
     class BS_CORE_EXPORT PixelData : public GpuResourceData
 	{
     public:
-    	PixelData() {}
+		PixelData();
 		~PixelData() {}
 
 		/**
 		 * Constructs a new object with an internal buffer capable of holding "extents" volume of pixels, where each pixel 
 		 * is of the specified pixel format. Extent offsets are also stored, but are not used internally.
 		 */
-		PixelData(const PixelVolume& extents, PixelFormat pixelFormat)
-			:mExtents(extents), mFormat(pixelFormat)
-		{
-			setConsecutive();
-		}
+		PixelData(const PixelVolume& extents, PixelFormat pixelFormat);
 
 		/**
 		 * Constructs a new object with an internal buffer capable of holding volume of pixels described by	provided width, 
 		 * height and depth, where each pixel is of the specified pixel format.
 		 */
-    	PixelData(UINT32 width, UINT32 height, UINT32 depth, PixelFormat pixelFormat)
-			: mExtents(0, 0, 0, width, height, depth), mFormat(pixelFormat)
-    	{
-    		setConsecutive();
-    	}
+		PixelData(UINT32 width, UINT32 height, UINT32 depth, PixelFormat pixelFormat);
 
 		PixelData(const PixelData& copy);
 		PixelData& operator=(const PixelData& rhs);
@@ -354,16 +346,6 @@ namespace bs
 		static SPtr<PixelData> create(UINT32 width, UINT32 height, UINT32 depth, PixelFormat pixelFormat);
 
 	private:
-		/**
-		 * Set the rowPitch and slicePitch so that the buffer is laid out consecutive in memory. Does not actually modify
-		 * the buffer itself.
-		 */
-		void setConsecutive()
-		{
-			mRowPitch = getWidth();
-			mSlicePitch = getWidth()*getHeight();
-		}
-
 		/**
 		 * Initializes the internal buffer with the provided set of colors. The array should be of width * height * depth 
 		 * size and mapped as such: arrayIdx = x + y * width + z * width * height.

Source/BansheeCore/Include/BsPixelUtil.h

@@ -87,6 +87,14 @@ namespace bs
 		static void getSizeForMipLevel(UINT32 width, UINT32 height, UINT32 depth, UINT32 mipLevel, 
 			UINT32& mipWidth, UINT32& mipHeight, UINT32& mipDepth);
 
+		/** 
+		 * Calculates row and depth pitch for a texture surface of the specified size and format. For most this will be
+		 * equal to their width & height, respectively. But some texture formats (especially compressed ones) might
+		 * require extra padding.
+		 */
+		static void getPitch(UINT32 width, UINT32 height, UINT32 depth, PixelFormat format,
+			UINT32& rowPitch, UINT32& depthPitch);
+
 		/**
 		 * Returns property flags for this pixel format.
 		 *
@@ -110,7 +118,7 @@ namespace bs
         static bool isNativeEndian(PixelFormat format);
 		
 		/**
-		 * Checks are the provided dimensions valid for the specified pixel format. Some formats (like DXT) require 
+		 * Checks are the provided dimensions valid for the specified pixel format. Some formats (like BC) require 
 		 * width/height to be multiples of four and some formats dont allow depth larger than 1.
 		 */
 		static bool isValidExtent(UINT32 width, UINT32 height, UINT32 depth, PixelFormat format);

Source/BansheeCore/Source/BsPixelData.cpp

@@ -8,6 +8,23 @@
 
 namespace bs
 {
+	PixelData::PixelData()
+		:mExtents(0, 0, 0, 0), mFormat(PF_UNKNOWN), mRowPitch(0), mSlicePitch(0)
+	{ }
+
+	PixelData::PixelData(const PixelVolume& extents, PixelFormat pixelFormat)
+		:mExtents(extents), mFormat(pixelFormat)
+	{
+		PixelUtil::getPitch(extents.getWidth(), extents.getHeight(), extents.getDepth(), pixelFormat, mRowPitch, 
+			mSlicePitch);
+	}
+
+	PixelData::PixelData(UINT32 width, UINT32 height, UINT32 depth, PixelFormat pixelFormat)
+		: mExtents(0, 0, 0, width, height, depth), mFormat(pixelFormat)
+	{
+		PixelUtil::getPitch(width, height, depth, pixelFormat, mRowPitch, mSlicePitch);
+	}
+
 	PixelData::PixelData(const PixelData& copy)
 		:GpuResourceData(copy)
 	{

Source/BansheeCore/Source/BsPixelUtil.cpp

@@ -982,14 +982,42 @@ namespace bs
 
 				default:
 					BS_EXCEPT(InvalidParametersException, "Invalid compressed pixel format");
+					return 0;
 			}
 		}
-		else
+
+		return width*height*depth*getNumElemBytes(format);
+	}
+
+	void PixelUtil::getPitch(UINT32 width, UINT32 height, UINT32 depth, PixelFormat format,
+						 UINT32& rowPitch, UINT32& depthPitch)
+	{
+		if (isCompressed(format))
 		{
-			return width*height*depth*getNumElemBytes(format);
+			switch (format)
+			{
+				// BC formats work by dividing the image into 4x4 blocks, then encoding each
+				// 4x4 block with a certain number of bytes. 
+			case PF_BC1:
+			case PF_BC1a:
+			case PF_BC4:
+			case PF_BC2:
+			case PF_BC3:
+			case PF_BC5:
+			case PF_BC6H:
+			case PF_BC7:
+				rowPitch = div(width + 3, 4).quot * 4;
+				depthPitch = div(height + 3, 4).quot * 4 * rowPitch;
+				return;
+
+			default:
+				BS_EXCEPT(InvalidParametersException, "Invalid compressed pixel format");
+				return;
+			}
 		}
 
-		return 0;
+		rowPitch = width;
+		depthPitch = width * height;
 	}
 
 	void PixelUtil::getSizeForMipLevel(UINT32 width, UINT32 height, UINT32 depth, UINT32 mipLevel,

Source/BansheeVulkanRenderAPI/Include/BsVulkanHardwareBuffer.h

@@ -16,7 +16,16 @@ namespace bs
 	class VulkanBuffer : public VulkanResource
 	{
 	public:
-		VulkanBuffer(VulkanResourceManager* owner, VkBuffer buffer, VkBufferView view, VkDeviceMemory memory);
+		/** 
+		 * @param[in]	owner		Manager that takes care of tracking and releasing of this object.
+		 * @param[in]	buffer		Actual low-level Vulkan buffer handle.
+		 * @param[in]	view		Optional handle to the buffer view.
+		 * @param[in]	memory		Memory mapped to the buffer.
+		 * @param[in]	rowPitch	If buffer maps to an image sub-resource, length of a single row (in elements).
+		 * @param[in]	slicePitch	If buffer maps to an image sub-resource, size of a single 2D surface (in elements).
+		 */
+		VulkanBuffer(VulkanResourceManager* owner, VkBuffer buffer, VkBufferView view, VkDeviceMemory memory, 
+			UINT32 rowPitch = 0, UINT32 slicePitch = 0);
 		~VulkanBuffer();
 
 		/** Returns the internal handle to the Vulkan object. */
@@ -25,6 +34,18 @@ namespace bs
 		/** Returns a buffer view that covers the entire buffer. */
 		VkBufferView getView() const { return mView; }
 
+		/**
+		 * If buffer represents an image sub-resource, this is the number of elements that separate one row of the 
+		 * sub-resource from another (if no padding, it is equal to image width).
+		 */
+		UINT32 getRowPitch() const { return mRowPitch; }
+
+		/**
+		 * If buffer represents an image sub-resource, this is the number of elements that separate one column of the
+		 * sub-resource from another (if no padding, it is equal to image height). Only relevant for 3D images.
+		 */
+		UINT32 getSliceHeight() const { return mSliceHeight; }
+
 		/** 
 		 * Returns a pointer to internal buffer memory. Must be followed by unmap(). Caller must ensure the buffer was
 		 * created in CPU readable memory, and that buffer isn't currently being written to by the GPU.
@@ -53,6 +74,9 @@ namespace bs
 		VkBuffer mBuffer;
 		VkBufferView mView;
 		VkDeviceMemory mMemory;
+
+		UINT32 mRowPitch;
+		UINT32 mSliceHeight;
 	};
 	
 	/**	Class containing common functionality for all Vulkan hardware buffers. */

Source/BansheeVulkanRenderAPI/Include/BsVulkanTexture.h

@@ -35,6 +35,20 @@ namespace bs
 		/** Returns an image view that covers the specified faces and mip maps of the texture. */
 		VkImageView getView(const TextureSurface& surface) const;
 
+		/** 
+		 * Returns a pointer to internal image memory for the specified sub-resource. Must be followed by unmap(). Caller
+		 * must ensure the image was created in CPU readable memory, and that image isn't currently being written to by the
+		 * GPU.
+		 * 
+		 * @param[in]	face		Index of the face to map.
+		 * @param[in]	mipLevel	Index of the mip level to map.
+		 * @param[in]	output		Output object containing the pointer to the sub-resource data.
+		 */
+		void map(UINT32 face, UINT32 mipLevel, PixelData& output) const;
+
+		/** Unmaps a buffer previously mapped with map(). */
+		void unmap();
+
 		/** 
 		 * Queues a command on the provided command buffer. The command copies the contents of the current image
 		 * subresource to the destination buffer. 
@@ -118,10 +132,14 @@ namespace bs
 		VulkanImage* createImage(VulkanDevice& device);
 
 		/** 
-		 * Creates a new buffer for the specified device, with enough space to hold the provided mip level of this
-		 * texture. 
+		 * Creates a staging buffer that can be used for texture transfer operations.
+		 * 
+		 * @param[in]	device		Device to create the buffer on.
+		 * @param[in]	pixelData	Object that describes the image sub-resource that will be in the buffer.
+		 * @param[in]	needsRead	True if we will be copying data from the buffer, false if just reading. True if both.
+		 * @return					Newly allocated buffer.
 		 */
-		VulkanBuffer* createStaging(VulkanDevice& device, UINT32 mipLevel, bool needsRead);
+		VulkanBuffer* createStaging(VulkanDevice& device, const PixelData& pixelData, bool needsRead);
 
 		VulkanImage* mImages[BS_MAX_DEVICES];
 		GpuDeviceFlags mDeviceMask;
@@ -132,6 +150,8 @@ namespace bs
 		UINT32 mMappedGlobalQueueIdx;
 		UINT32 mMappedMip;
 		UINT32 mMappedFace;
+		UINT32 mMappedRowPitch;
+		UINT32 mMappedSlicePitch;
 		GpuLockOptions mMappedLockOptions;
 
 		VkImageCreateInfo mImageCI;

Source/BansheeVulkanRenderAPI/Source/BsVulkanHardwareBuffer.cpp

@@ -10,8 +10,10 @@
 
 namespace bs
 {
-	VulkanBuffer::VulkanBuffer(VulkanResourceManager* owner, VkBuffer buffer, VkBufferView view, VkDeviceMemory memory)
-		:VulkanResource(owner, false), mBuffer(buffer), mView(view), mMemory(memory)
+	VulkanBuffer::VulkanBuffer(VulkanResourceManager* owner, VkBuffer buffer, VkBufferView view, VkDeviceMemory memory,
+							   UINT32 rowPitch, UINT32 slicePitch)
+		: VulkanResource(owner, false), mBuffer(buffer), mView(view), mMemory(memory), mRowPitch(rowPitch)
+		, mSliceHeight(slicePitch / rowPitch)
 	{
 
 	}
@@ -60,8 +62,8 @@ namespace bs
 		const VkImageSubresourceLayers& range, VkImageLayout layout)
 	{
 		VkBufferImageCopy region;
-		region.bufferRowLength = 0;
-		region.bufferImageHeight = 0;
+		region.bufferRowLength = mRowPitch;
+		region.bufferImageHeight = mSliceHeight;
 		region.bufferOffset = 0;
 		region.imageOffset.x = 0;
 		region.imageOffset.y = 0;

Source/BansheeVulkanRenderAPI/Source/BsVulkanTexture.cpp

@@ -125,12 +125,46 @@ namespace bs
 		return view;
 	}
 
+	void VulkanImage::map(UINT32 face, UINT32 mipLevel, PixelData& output) const
+	{
+		VulkanDevice& device = mOwner->getDevice();
+
+		VkImageSubresource range;
+		range.mipLevel = mipLevel;
+		range.arrayLayer = face;
+
+		if (mImageViewCI.subresourceRange.aspectMask == VK_IMAGE_ASPECT_COLOR_BIT)
+			range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
+		else // Depth stencil, but we only map depth
+			range.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;
+
+		VkSubresourceLayout layout;
+		vkGetImageSubresourceLayout(device.getLogical(), mImage, &range, &layout);
+
+		assert(layout.size == output.getSize());
+		output.setRowPitch(layout.rowPitch);
+		output.setSlicePitch(layout.depthPitch);
+
+		UINT8* data;
+		VkResult result = vkMapMemory(device.getLogical(), mMemory, layout.offset, layout.size, 0, (void**)&data);
+		assert(result == VK_SUCCESS);
+
+		output.setExternalBuffer(data);
+	}
+
+	void VulkanImage::unmap()
+	{
+		VulkanDevice& device = mOwner->getDevice();
+
+		vkUnmapMemory(device.getLogical(), mMemory);
+	}
+
 	void VulkanImage::copy(VulkanTransferBuffer* cb, VulkanBuffer* destination, const VkExtent3D& extent,
-							const VkImageSubresourceLayers& range, VkImageLayout layout)
+						   const VkImageSubresourceLayers& range, VkImageLayout layout)
 	{
 		VkBufferImageCopy region;
-		region.bufferRowLength = 0;
-		region.bufferImageHeight = 0;
+		region.bufferRowLength = destination->getRowPitch();
+		region.bufferImageHeight = destination->getSliceHeight();
 		region.bufferOffset = 0;
 		region.imageOffset.x = 0;
 		region.imageOffset.y = 0;
@@ -145,7 +179,8 @@ namespace bs
 		GpuDeviceFlags deviceMask)
 		: TextureCore(desc, initialData, deviceMask), mImages(), mDeviceMask(deviceMask), mAccessFlags(0)
 		, mStagingBuffer(nullptr), mMappedDeviceIdx(-1), mMappedGlobalQueueIdx(-1), mMappedMip(0), mMappedFace(0)
-		, mMappedLockOptions(GBL_WRITE_ONLY), mDirectlyMappable(false), mSupportsGPUWrites(false), mIsMapped(false)
+		, mMappedRowPitch(false), mMappedSlicePitch(false), mMappedLockOptions(GBL_WRITE_ONLY), mDirectlyMappable(false)
+		, mSupportsGPUWrites(false), mIsMapped(false)
 	{
 		
 	}
@@ -290,21 +325,13 @@ namespace bs
 		return device.getResourceManager().create<VulkanImage>(image, memory, mImageCI.initialLayout, getProperties());
 	}
 
-	VulkanBuffer* VulkanTextureCore::createStaging(VulkanDevice& device, UINT32 mipLevel, bool readable)
+	VulkanBuffer* VulkanTextureCore::createStaging(VulkanDevice& device, const PixelData& pixelData, bool readable)
 	{
-		const TextureProperties& props = getProperties();
-
-		UINT32 mipWidth, mipHeight, mipDepth;
-		PixelUtil::getSizeForMipLevel(props.getWidth(), props.getHeight(), props.getDepth(), mipLevel, mipWidth, mipHeight, 
-			mipDepth);
-
-		UINT32 mipLevelSize = PixelUtil::getMemorySize(mipWidth, mipHeight, mipDepth, props.getFormat());
-
 		VkBufferCreateInfo bufferCI;
 		bufferCI.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
 		bufferCI.pNext = nullptr;
 		bufferCI.flags = 0;
-		bufferCI.size = mipLevelSize;
+		bufferCI.size = pixelData.getSize();
 		bufferCI.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
 		bufferCI.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
 		bufferCI.queueFamilyIndexCount = 0;
@@ -328,7 +355,8 @@ namespace bs
 		result = vkBindBufferMemory(vkDevice, buffer, memory, 0);
 		assert(result == VK_SUCCESS);
 
-		return device.getResourceManager().create<VulkanBuffer>(buffer, VK_NULL_HANDLE, memory);
+		return device.getResourceManager().create<VulkanBuffer>(buffer, VK_NULL_HANDLE, memory, 
+			pixelData.getRowPitch(), pixelData.getSlicePitch());
 	}
 
 	VkImageView VulkanTextureCore::getView(UINT32 deviceIdx) const
@@ -356,7 +384,9 @@ namespace bs
 	PixelData VulkanTextureCore::lockImpl(GpuLockOptions options, UINT32 mipLevel, UINT32 face, UINT32 deviceIdx,
 										  UINT32 queueIdx)
 	{
-		if (mProperties.getNumSamples() > 1)
+		const TextureProperties& props = getProperties();
+
+		if (props.getNumSamples() > 1)
 		{
 			LOGERR("Multisampled textures cannot be accessed from the CPU directly.");
 			return PixelData();
@@ -374,11 +404,11 @@ namespace bs
 		}
 #endif
 
-		UINT32 mipWidth = std::max(1u, mProperties.getWidth() >> mipLevel);
-		UINT32 mipHeight = std::max(1u, mProperties.getHeight() >> mipLevel);
-		UINT32 mipDepth = std::max(1u, mProperties.getDepth() >> mipLevel);
+		UINT32 mipWidth = std::max(1u, props.getWidth() >> mipLevel);
+		UINT32 mipHeight = std::max(1u, props.getHeight() >> mipLevel);
+		UINT32 mipDepth = std::max(1u, props.getDepth() >> mipLevel);
 
-		PixelData lockedArea(mipWidth, mipHeight, mipDepth, mProperties.getFormat());
+		PixelData lockedArea(mipWidth, mipHeight, mipDepth, props.getFormat());
 
 		VulkanImage* image = mImages[deviceIdx];
 
@@ -418,16 +448,25 @@ namespace bs
 
 			// We're safe to map directly since GPU isn't using the subresource
 			if (!isUsedOnGPU)
-				return image->map(face, mipLevel);
+			{
+				image->map(face, mipLevel, lockedArea);
+				return lockedArea;
+			}
 
 			// Caller guarantees he won't touch the same data as the GPU, so just map even though the GPU is using the
 			// subresource
 			if (options == GBL_WRITE_ONLY_NO_OVERWRITE)
-				return image->map(face, mipLevel);
+			{
+				image->map(face, mipLevel, lockedArea);
+				return lockedArea;
+			}
 
 			// No GPU writes are are supported and we're only reading, so no need to wait on anything
 			if (options == GBL_READ_ONLY)
-				return image->map(face, mipLevel);
+			{
+				image->map(face, mipLevel, lockedArea);
+				return lockedArea;
+			}
 
 			// Caller doesn't care about buffer contents, so just discard the existing buffer and create a new one
 			if (options == GBL_WRITE_ONLY_DISCARD)
@@ -453,7 +492,8 @@ namespace bs
 				// Submit the command buffer and wait until it finishes
 				transferCB->flush(true);
 
-				return image->map(face, mipLevel);
+				image->map(face, mipLevel, lockedArea);
+				return lockedArea;
 			}
 
 			// Otherwise, we're doing write only, in which case it's best to use the staging buffer to avoid waiting
@@ -463,7 +503,7 @@ namespace bs
 		bool needRead = options == GBL_READ_WRITE || options == GBL_READ_ONLY;
 
 		// Allocate a staging buffer
-		mStagingBuffer = createStaging(device, mipLevel, needRead);
+		mStagingBuffer = createStaging(device, lockedArea, needRead);
 
 		if (needRead) // If reading, we need to copy the current contents of the image to the staging buffer
 		{
@@ -477,8 +517,6 @@ namespace bs
 				transferCB->appendMask(writeUseMask);
 			}
 
-			const TextureProperties& props = getProperties();
-
 			VkImageSubresourceRange range;
 			if ((props.getUsage() & TU_DEPTHSTENCIL) != 0)
 				range.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT;
@@ -532,7 +570,10 @@ namespace bs
 			transferCB->flush(true);
 		}
 
-		return mStagingBuffer->map(offset, length);
+		UINT8* data = mStagingBuffer->map(0, lockedArea.getSize());
+		lockedArea.setExternalBuffer(data);
+
+		return lockedArea;
 	}
 
 	void VulkanTextureCore::unlockImpl()