Vulkan: Handling layout transitions for textures used both as framebuffer attachments and shader inputs

Source/BansheeVulkanRenderAPI/Include/BsVulkanCommandBuffer.h

@@ -160,7 +160,7 @@ namespace bs
 		 * device when the command buffer is submitted.
 		 */
 		void registerResource(VulkanImage* res, VkAccessFlags accessFlags, VkImageLayout layout, 
-			const VkImageSubresourceRange& range, VulkanUseFlags flags);
+			VulkanUseFlags flags, bool isFBAttachment = false);
 
 		/** 
 		 * Lets the command buffer know that the provided image resource has been queued on it, and will be used by the
@@ -254,9 +254,16 @@ namespace bs
 		struct ImageInfo
 		{
 			VkAccessFlags accessFlags;
-			VkImageLayout layout;
 			VkImageSubresourceRange range;
 			ResourceUseHandle useHandle;
+
+			// Only relevant for layout transitions
+			VkImageLayout currentLayout;
+			VkImageLayout requiredLayout;
+			VkImageLayout finalLayout;
+
+			bool isFBAttachment : 1;
+			bool isShaderInput : 1;
 		};
 
 		/** Checks if all the prerequisites for rendering have been made (e.g. render target and pipeline state are set. */
@@ -293,8 +300,9 @@ namespace bs
 		RenderSurfaceMask mRenderTargetLoadMask;
 
 		UnorderedMap<VulkanResource*, ResourceUseHandle> mResources;
-		UnorderedMap<VulkanResource*, ImageInfo> mImages;
+		UnorderedMap<VulkanResource*, UINT32> mImages;
 		UnorderedMap<VulkanResource*, BufferInfo> mBuffers;
+		Vector<ImageInfo> mImageInfos;
 		UINT32 mGlobalQueueIdx;
 
 		SPtr<VulkanGraphicsPipelineStateCore> mGraphicsPipeline;
@@ -317,6 +325,8 @@ namespace bs
 		VkDeviceSize mVertexBufferOffsetsTemp[BS_MAX_BOUND_VERTEX_BUFFERS];
 		VkDescriptorSet* mDescriptorSetsTemp;
 		UnorderedMap<UINT32, TransitionInfo> mTransitionInfoTemp;
+		Vector<VkImageMemoryBarrier> mLayoutTransitionBarriersTemp;
+		UnorderedMap<VulkanImage*, UINT32> mQueuedLayoutTransitions;
 	};
 
 	/** CommandBuffer implementation for Vulkan. */
@@ -354,9 +364,6 @@ namespace bs
 		VulkanDevice& mDevice;
 		VulkanQueue* mQueue;
 		UINT32 mIdMask;
-
-		VkSemaphore mSemaphoresTemp[BS_MAX_UNIQUE_QUEUES];
-		UnorderedMap<UINT32, TransitionInfo> mTransitionInfoTemp;
 	};
 
 	/** @} */

Source/BansheeVulkanRenderAPI/Include/BsVulkanFramebuffer.h

@@ -52,6 +52,14 @@ namespace bs
 		bool offscreen;
 	};
 
+	/** Information about a single framebuffer attachment. */
+	struct VulkanFramebufferAttachment
+	{
+		VulkanImage* image = nullptr;
+		UINT32 baseLayer = 0;
+		VkImageLayout finalLayout = VK_IMAGE_LAYOUT_UNDEFINED;
+	};
+
 	/** Vulkan frame buffer containing one or multiple color surfaces, and an optional depth surface. */
 	class VulkanFramebuffer : public VulkanResource
 	{
@@ -89,17 +97,11 @@ namespace bs
 		 */
 		UINT32 getNumLayers() const { return mNumLayers; }
 
-		/** Returns the image representing the color attachment at the provided index. */
-		VulkanImage* getColorImage(UINT32 colorIdx) const { return mColorImages[colorIdx]; }
-
-		/** Returns the image representing the depth/stencil attachment, if one exists. */
-		VulkanImage* getDepthStencilImage() const { return mDepthStencilImage; }
-
-		/** Returns the initial layer of the color texture surface in which to start rendering. */
-		UINT32 getColorBaseLayer(UINT32 colorIdx) const { return mColorBaseLayers[colorIdx]; }
+		/** Returns information about a color attachment at the specified index. */
+		const VulkanFramebufferAttachment& getColorAttachment(UINT32 colorIdx) const { return mColorAttachments[colorIdx]; }
 
-		/** Returns the initial layer of the depth-stencil texture surface in which to start rendering. */
-		UINT32 getDepthStencilBaseLayer() const { return mDepthBaseLayer; }
+		/** Returns information about a depth-stencil attachment. */
+		const VulkanFramebufferAttachment& getDepthStencilAttachment() const { return mDepthStencilAttachment; }
 
 		/** Gets the total number of frame-buffer attachments, including both color and depth. */
 		UINT32 getNumAttachments() const { return mNumAttachments; }
@@ -152,10 +154,8 @@ namespace bs
 		UINT32 mNumAttachments;
 		UINT32 mNumColorAttachments;
 		UINT32 mNumLayers;
-		VulkanImage* mColorImages[BS_MAX_MULTIPLE_RENDER_TARGETS];
-		VulkanImage* mDepthStencilImage;
-		UINT32 mColorBaseLayers[BS_MAX_MULTIPLE_RENDER_TARGETS];
-		UINT32 mDepthBaseLayer;
+		VulkanFramebufferAttachment mColorAttachments[BS_MAX_MULTIPLE_RENDER_TARGETS];
+		VulkanFramebufferAttachment mDepthStencilAttachment;
 		bool mHasDepth;
 		VkSampleCountFlagBits mSampleFlags;
 

Source/BansheeVulkanRenderAPI/Include/BsVulkanTexture.h

@@ -66,6 +66,9 @@ namespace bs
 
 		/** Returns an image view that covers the specified faces and mip maps of the texture. */
 		VkImageView getView(const TextureSurface& surface) const;
+		
+		/** Retrieves a subresource range covering all the sub-resources of the image. */
+		VkImageSubresourceRange getRange() const;
 
 		/** 
 		 * Retrieves a separate resource for a specific image face & mip level. This allows the caller to track subresource

Source/BansheeVulkanRenderAPI/Source/BsVulkanCommandBuffer.cpp

@@ -170,7 +170,10 @@ namespace bs
 
 			for (auto& entry : mImages)
 			{
-				ResourceUseHandle& useHandle = entry.second.useHandle;
+				UINT32 imageInfoIdx = entry.second;
+				ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+				ResourceUseHandle& useHandle = imageInfo.useHandle;
 				assert(useHandle.used);
 
 				entry.first->notifyUnbound();
@@ -233,11 +236,97 @@ namespace bs
 			return;
 		}
 
+		// Perform any queued layout transitions
+		auto createLayoutTransitionBarrier = [&](VulkanImage* image, ImageInfo& imageInfo)
+		{
+			VkAccessFlags srcAccessFlags;
+			if (imageInfo.currentLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
+				srcAccessFlags = VK_ACCESS_SHADER_READ_BIT;
+			else if (imageInfo.currentLayout == VK_IMAGE_LAYOUT_GENERAL)
+			{
+				srcAccessFlags = VK_ACCESS_SHADER_READ_BIT;
+				srcAccessFlags |= VK_ACCESS_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT;
+			}
+			else if(imageInfo.currentLayout == VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL)
+				srcAccessFlags = VK_ACCESS_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
+			else if (imageInfo.currentLayout == VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL)
+				srcAccessFlags = VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
+			else
+			{
+				srcAccessFlags = VK_ACCESS_SHADER_READ_BIT;
+				LOGWRN("Unsupported source layout for a framebuffer attachment.");
+			}
+
+			mLayoutTransitionBarriersTemp.push_back(VkImageMemoryBarrier());
+			VkImageMemoryBarrier& barrier = mLayoutTransitionBarriersTemp.back();
+			barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
+			barrier.pNext = nullptr;
+			barrier.srcAccessMask = srcAccessFlags;
+			barrier.dstAccessMask = imageInfo.accessFlags;
+			barrier.srcQueueFamilyIndex = mQueueFamily;
+			barrier.dstQueueFamilyIndex = mQueueFamily;
+			barrier.oldLayout = imageInfo.currentLayout;
+			barrier.newLayout = imageInfo.requiredLayout;
+			barrier.image = image->getHandle();
+			barrier.subresourceRange = imageInfo.range;
+
+			imageInfo.currentLayout = imageInfo.requiredLayout;
+		};
+
+		for (auto& entry : mQueuedLayoutTransitions)
+		{
+			UINT32 imageInfoIdx = entry.second;
+			ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+			createLayoutTransitionBarrier(entry.first, imageInfo);
+		}
+
+		mQueuedLayoutTransitions.clear();
+
+		vkCmdPipelineBarrier(mCmdBuffer,
+							 VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, // Note: VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT might be more correct here, according to the spec
+							 VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
+							 0, 0, nullptr,
+							 0, nullptr,
+							 (UINT32)mLayoutTransitionBarriersTemp.size(), mLayoutTransitionBarriersTemp.data());
+
+		mLayoutTransitionBarriersTemp.clear();
+
+		// Check if any frame-buffer attachments are also used as shader inputs, in which case we make them read-only
+		RenderSurfaceMask readMask = RT_NONE;
+
+		UINT32 numColorAttachments = mFramebuffer->getNumColorAttachments();
+		for(UINT32 i = 0; i < numColorAttachments; i++)
+		{
+			VulkanImage* image = mFramebuffer->getColorAttachment(i).image;
+
+			UINT32 imageInfoIdx = mImages[image];
+			ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+			bool readOnly = imageInfo.isShaderInput;
+
+			if(readOnly)
+				readMask.set((RenderSurfaceMaskBits)(1 << i));
+		}
+
+		if(mFramebuffer->hasDepthAttachment())
+		{
+			VulkanImage* image = mFramebuffer->getDepthStencilAttachment().image;
+
+			UINT32 imageInfoIdx = mImages[image];
+			ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+			bool readOnly = imageInfo.isShaderInput;
+
+			if (readOnly)
+				readMask.set(RT_DEPTH);
+		}
+
 		VkRenderPassBeginInfo renderPassBeginInfo;
 		renderPassBeginInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
 		renderPassBeginInfo.pNext = nullptr;
-		renderPassBeginInfo.framebuffer = mFramebuffer->getFramebuffer(mRenderTargetLoadMask, RT_NONE);
-		renderPassBeginInfo.renderPass = mFramebuffer->getRenderPass(mRenderTargetLoadMask, RT_NONE);
+		renderPassBeginInfo.framebuffer = mFramebuffer->getFramebuffer(mRenderTargetLoadMask, readMask);
+		renderPassBeginInfo.renderPass = mFramebuffer->getRenderPass(mRenderTargetLoadMask, readMask);
 		renderPassBeginInfo.renderArea.offset.x = 0;
 		renderPassBeginInfo.renderArea.offset.y = 0;
 		renderPassBeginInfo.renderArea.extent.width = mRenderTargetWidth;
@@ -293,11 +382,12 @@ namespace bs
 		for (auto& entry : mImages)
 		{
 			VulkanImage* resource = static_cast<VulkanImage*>(entry.first);
+			ImageInfo& imageInfo = mImageInfos[entry.second];
 
 			UINT32 currentQueueFamily = resource->getQueueFamily();
 			bool queueMismatch = resource->isExclusive() && currentQueueFamily != mQueueFamily;
 
-			if (queueMismatch || resource->getLayout() != entry.second.layout)
+			if (queueMismatch || imageInfo.currentLayout != imageInfo.requiredLayout)
 			{
 				Vector<VkImageMemoryBarrier>& barriers = mTransitionInfoTemp[currentQueueFamily].imageBarriers;
 
@@ -305,16 +395,17 @@ namespace bs
 				VkImageMemoryBarrier& barrier = barriers.back();
 				barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
 				barrier.pNext = nullptr;
-				barrier.srcAccessMask = entry.second.accessFlags;
-				barrier.dstAccessMask = entry.second.accessFlags;
+				barrier.srcAccessMask = imageInfo.accessFlags;
+				barrier.dstAccessMask = imageInfo.accessFlags;
 				barrier.srcQueueFamilyIndex = currentQueueFamily;
 				barrier.dstQueueFamilyIndex = mQueueFamily;
-				barrier.oldLayout = resource->getLayout();
-				barrier.newLayout = entry.second.layout;
+				barrier.oldLayout = imageInfo.currentLayout;
+				barrier.newLayout = imageInfo.requiredLayout;
 				barrier.image = resource->getHandle();
-				barrier.subresourceRange = entry.second.range;
+				barrier.subresourceRange = imageInfo.range;
 
-				resource->setLayout(entry.second.layout);
+				imageInfo.currentLayout = imageInfo.requiredLayout;
+				resource->setLayout(imageInfo.finalLayout);
 			}
 		}
 
@@ -441,7 +532,10 @@ namespace bs
 
 		for (auto& entry : mImages)
 		{
-			ResourceUseHandle& useHandle = entry.second.useHandle;
+			UINT32 imageInfoIdx = entry.second;
+			ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+			ResourceUseHandle& useHandle = imageInfo.useHandle;
 			assert(!useHandle.used);
 
 			useHandle.used = true;
@@ -477,6 +571,7 @@ namespace bs
 		mCmpPipelineRequiresBind = true;
 		mFramebuffer = nullptr;
 		mDescriptorSetsBindState = DescriptorSetBindFlag::Graphics | DescriptorSetBindFlag::Compute;
+		mQueuedLayoutTransitions.clear();
 	}
 
 	void VulkanCmdBuffer::refreshFenceStatus()
@@ -508,7 +603,10 @@ namespace bs
 
 				for (auto& entry : mImages)
 				{
-					ResourceUseHandle& useHandle = entry.second.useHandle;
+					UINT32 imageInfoIdx = entry.second;
+					ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+					ResourceUseHandle& useHandle = imageInfo.useHandle;
 					assert(useHandle.used);
 
 					entry.first->notifyDone(mGlobalQueueIdx, useHandle.flags);
@@ -525,6 +623,7 @@ namespace bs
 				mResources.clear();
 				mImages.clear();
 				mBuffers.clear();
+				mImageInfos.clear();
 			}
 		}
 		else
@@ -560,8 +659,6 @@ namespace bs
 			mRenderTargetHeight = rt->getProperties().getHeight();
 			mRenderTargetDepthReadOnly = readOnlyDepthStencil;
 			mRenderTargetLoadMask = loadMask;
-
-			registerResource(mFramebuffer, VulkanUseFlag::Write);
 		}
 
 		// If anything changed
@@ -570,6 +667,19 @@ namespace bs
 			if (isInRenderPass())
 				endRenderPass();
 
+			// Reset flags that signal image usage
+			for (auto& entry : mImages)
+			{
+				UINT32 imageInfoIdx = entry.second;
+				ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+				imageInfo.isFBAttachment = false;
+				imageInfo.isShaderInput = false;
+			}
+
+			setGpuParams(nullptr);
+
+			registerResource(mFramebuffer, VulkanUseFlag::Write);
 			mGfxPipelineRequiresBind = true;
 		}
 	}
@@ -601,7 +711,7 @@ namespace bs
 				colorValue.float32[2] = color.b;
 				colorValue.float32[3] = color.a;
 
-				UINT32 curBaseLayer = mFramebuffer->getColorBaseLayer(i);
+				UINT32 curBaseLayer = mFramebuffer->getColorAttachment(i).baseLayer;
 				if (attachmentIdx == 0)
 					baseLayer = curBaseLayer;
 				else
@@ -638,7 +748,7 @@ namespace bs
 
 				attachments[attachmentIdx].colorAttachment = 0;
 
-				UINT32 curBaseLayer = mFramebuffer->getDepthStencilBaseLayer();
+				UINT32 curBaseLayer = mFramebuffer->getDepthStencilAttachment().baseLayer;
 				if (attachmentIdx == 0)
 					baseLayer = curBaseLayer;
 				else
@@ -839,6 +949,36 @@ namespace bs
 		if (pipeline == nullptr)
 			return false;
 
+		// Check that pipeline matches the read-only state of any framebuffer attachments
+		UINT32 numColorAttachments = mFramebuffer->getNumColorAttachments();
+		for (UINT32 i = 0; i < numColorAttachments; i++)
+		{
+			VulkanImage* image = mFramebuffer->getColorAttachment(i).image;
+
+			UINT32 imageInfoIdx = mImages[image];
+			ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+			if (imageInfo.isShaderInput && !pipeline->isColorReadOnly(i))
+			{
+				LOGWRN("Framebuffer attachment also used as a shader input, but color writes aren't disabled. This will"
+					" result in undefined behavior.");
+			}
+		}
+
+		if (mFramebuffer->hasDepthAttachment())
+		{
+			VulkanImage* image = mFramebuffer->getDepthStencilAttachment().image;
+
+			UINT32 imageInfoIdx = mImages[image];
+			ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
+
+			if (imageInfo.isShaderInput && !pipeline->isDepthStencilReadOnly())
+			{
+				LOGWRN("Framebuffer attachment also used as a shader input, but depth/stencil writes aren't disabled. "
+					"This will result in undefined behavior.");
+			}
+		}
+
 		mGraphicsPipeline->registerPipelineResources(this);
 		registerResource(pipeline, VulkanUseFlag::Read);
 
@@ -1010,31 +1150,91 @@ namespace bs
 	}
 
 	void VulkanCmdBuffer::registerResource(VulkanImage* res, VkAccessFlags accessFlags, VkImageLayout layout, 
-		const VkImageSubresourceRange& range, VulkanUseFlags flags)
+		VulkanUseFlags flags, bool isFBAttachment)
 	{
-		auto insertResult = mImages.insert(std::make_pair(res, ImageInfo()));
+		// Note: I currently always perform pipeline barriers (layout transitions and similar), over the entire image.
+		//       In the case of render and storage images, the case is often that only a specific subresource requires
+		//       it. However this makes grouping and tracking of current image layouts much more difficult.
+		//       If this is ever requires we'll need to track image layout per-subresource instead per-image, and we
+		//       might also need a smart way to group layout transitions for multiple sub-resources on the same image.
+
+		VkImageSubresourceRange range = res->getRange();
+		UINT32 nextImageInfoIdx = (UINT32)mImageInfos.size();
+
+		auto insertResult = mImages.insert(std::make_pair(res, nextImageInfoIdx));
 		if (insertResult.second) // New element
 		{
-			ImageInfo& imageInfo = insertResult.first->second;
+			UINT32 imageInfoIdx = insertResult.first->second;
+			mImageInfos.push_back(ImageInfo());
+
+			ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
 			imageInfo.accessFlags = accessFlags;
-			imageInfo.layout = layout;
+			imageInfo.currentLayout = res->getLayout();
+			imageInfo.requiredLayout = layout;
+			imageInfo.finalLayout = layout;
 			imageInfo.range = range;
+			imageInfo.isFBAttachment = isFBAttachment;
+			imageInfo.isShaderInput = !isFBAttachment;
 
 			imageInfo.useHandle.used = false;
 			imageInfo.useHandle.flags = flags;
 
 			res->notifyBound();
+
+			if (imageInfo.currentLayout != imageInfo.requiredLayout)
+				mQueuedLayoutTransitions[res] = imageInfoIdx;
 		}
 		else // Existing element
 		{
-			ImageInfo& imageInfo = insertResult.first->second;
+			UINT32 imageInfoIdx = insertResult.first->second;
+			ImageInfo& imageInfo = mImageInfos[imageInfoIdx];
 
 			assert(!imageInfo.useHandle.used);
 			imageInfo.useHandle.flags |= flags;
 
-			assert(imageInfo.layout == layout && "Cannot bind the same image with two different layouts on the same command buffer.");
 			imageInfo.accessFlags |= accessFlags;
-			imageInfo.range = range;
+
+			// Check if the same image is used with different layouts, in which case we need to transfer to the general
+			// layout
+			if (imageInfo.requiredLayout != layout)
+				imageInfo.requiredLayout = VK_IMAGE_LAYOUT_GENERAL;
+
+			// If attached to FB, then the final layout is set by the FB (provided as layout param here), otherwise its
+			// the same as required layout
+			if(!isFBAttachment && !imageInfo.isFBAttachment)
+				imageInfo.finalLayout = imageInfo.requiredLayout;
+			else
+			{
+				if (isFBAttachment)
+					imageInfo.finalLayout = layout;
+			}
+
+			if (imageInfo.currentLayout != imageInfo.requiredLayout)
+				mQueuedLayoutTransitions[res] = imageInfoIdx;
+
+			// If a FB attachment was just bound as a shader input, we might need to restart the render pass with a FB
+			// attachment that supports read-only attachments using the GENERAL layout
+			bool requiresReadOnlyFB = false;
+			if (isFBAttachment)
+			{
+				if (!imageInfo.isFBAttachment)
+				{
+					imageInfo.isFBAttachment = true;
+					requiresReadOnlyFB = imageInfo.isShaderInput;
+				}
+			}
+			else
+			{
+				if (!imageInfo.isShaderInput)
+				{
+					imageInfo.isShaderInput = true;
+					requiresReadOnlyFB = imageInfo.isFBAttachment;
+				}
+			}
+
+			// If we need to switch frame-buffers, end current render pass
+			if (requiresReadOnlyFB && isInRenderPass())
+				endRenderPass();
 		}
 
 		// Register any sub-resources
@@ -1093,18 +1293,20 @@ namespace bs
 		}
 
 		// Register any sub-resources
-		// (Purposely don't register them as images, as we will handle any layout transitions manually)
 		UINT32 numColorAttachments = res->getNumColorAttachments();
 		for (UINT32 i = 0; i < numColorAttachments; i++)
 		{
-			VulkanImage* image = res->getColorImage(i);
-			registerResource(image, VulkanUseFlag::Write);
+			const VulkanFramebufferAttachment& attachment = res->getColorAttachment(i);
+			registerResource(attachment.image, VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_COLOR_ATTACHMENT_READ_BIT,
+							 attachment.finalLayout, VulkanUseFlag::Write, true);
 		}
 
 		if(res->hasDepthAttachment())
 		{
-			VulkanImage* image = res->getDepthStencilImage();
-			registerResource(image, VulkanUseFlag::Write);
+			const VulkanFramebufferAttachment& attachment = res->getDepthStencilAttachment();
+			registerResource(attachment.image,
+							 VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT,
+							 attachment.finalLayout, VulkanUseFlag::Write, true);
 		}
 	}
 

Source/BansheeVulkanRenderAPI/Source/BsVulkanFramebuffer.cpp

@@ -30,8 +30,7 @@ namespace bs
 
 	VulkanFramebuffer::VulkanFramebuffer(VulkanResourceManager* owner, const VULKAN_FRAMEBUFFER_DESC& desc)
 		: VulkanResource(owner, false), mNumAttachments(0), mNumColorAttachments(0), mNumLayers(desc.layers)
-		, mColorImages(), mDepthStencilImage(nullptr), mColorBaseLayers(), mDepthBaseLayer(0), mHasDepth(false)
-		, mSampleFlags(VK_SAMPLE_COUNT_1_BIT)
+		, mColorAttachments(), mDepthStencilAttachment(), mHasDepth(false), mSampleFlags(VK_SAMPLE_COUNT_1_BIT)
 	{
 		mId = sNextValidId++;
 
@@ -59,8 +58,9 @@ namespace bs
 			else
 				attachmentDesc.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
 
-			mColorBaseLayers[attachmentIdx] = desc.color[i].baseLayer;
-			mColorImages[attachmentIdx] = desc.color[i].image;
+			mColorAttachments[attachmentIdx].baseLayer = desc.color[i].baseLayer;
+			mColorAttachments[attachmentIdx].image = desc.color[i].image;
+			mColorAttachments[attachmentIdx].finalLayout = attachmentDesc.finalLayout;
 
 			VkAttachmentReference& ref = mColorReferences[attachmentIdx];
 			ref.attachment = attachmentIdx;
@@ -86,8 +86,9 @@ namespace bs
 			attachmentDesc.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
 			attachmentDesc.finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
 
-			mDepthBaseLayer = desc.depth.baseLayer;
-			mDepthStencilImage = desc.depth.image;
+			mDepthStencilAttachment.baseLayer = desc.depth.baseLayer;
+			mDepthStencilAttachment.image = desc.depth.image;
+			mDepthStencilAttachment.finalLayout = attachmentDesc.finalLayout;
 
 			VkAttachmentReference& ref = mDepthReference;
 			ref.attachment = attachmentIdx;

Source/BansheeVulkanRenderAPI/Source/BsVulkanGpuParams.cpp

@@ -372,18 +372,7 @@ namespace bs
 			VkAccessFlags accessFlags = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;
 			VulkanUseFlags useFlags = VulkanUseFlag::Read | VulkanUseFlag::Write;
 
-			const TextureSurface& surface = mLoadStoreSurfaces[i];
-			VkImageSubresourceRange range;
-			range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
-
-			// Note: Currently layout transitions are only supported for entire images, not individual subresources, so
-			// make sure to bind all faces/mipmaps
-			range.baseArrayLayer = 0;
-			range.layerCount = props.getNumFaces();
-			range.baseMipLevel = 0;
-			range.levelCount = props.getNumMipmaps();
-
-			buffer.registerResource(resource, accessFlags, VK_IMAGE_LAYOUT_GENERAL, range, useFlags);
+			buffer.registerResource(resource, accessFlags, VK_IMAGE_LAYOUT_GENERAL, useFlags);
 		}
 
 		for (UINT32 i = 0; i < numTextures; i++)
@@ -399,17 +388,6 @@ namespace bs
 
 			const TextureProperties& props = element->getProperties();
 
-			bool isDepthStencil = (props.getUsage() & TU_DEPTHSTENCIL) != 0;
-
-			VkImageSubresourceRange range;
-			range.aspectMask = isDepthStencil ? VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT
-				: VK_IMAGE_ASPECT_COLOR_BIT;
-
-			range.baseArrayLayer = 0;
-			range.layerCount = props.getNumFaces();
-			range.baseMipLevel = 0;
-			range.levelCount = props.getNumMipmaps();
-
 			VkImageLayout layout;
 
 			// Keep dynamic textures in general layout, so they can be easily mapped by CPU
@@ -418,8 +396,7 @@ namespace bs
 			else
 				layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
 
-			buffer.registerResource(resource, VK_ACCESS_SHADER_READ_BIT, layout,
-				range, VulkanUseFlag::Read);
+			buffer.registerResource(resource, VK_ACCESS_SHADER_READ_BIT, layout, VulkanUseFlag::Read);
 		}
 
 		// Acquire sets as needed, and updated their contents if dirty

Source/BansheeVulkanRenderAPI/Source/BsVulkanTexture.cpp

@@ -161,6 +161,18 @@ namespace bs
 		return view;
 	}
 
+	VkImageSubresourceRange VulkanImage::getRange() const
+	{
+		VkImageSubresourceRange range;
+		range.aspectMask = mImageViewCI.subresourceRange.aspectMask;
+		range.baseArrayLayer = 0;
+		range.layerCount = mNumFaces;
+		range.baseMipLevel = 0;
+		range.levelCount = mNumMipLevels;
+
+		return range;
+	}
+
 	VulkanImageSubresource* VulkanImage::getSubresource(UINT32 face, UINT32 mipLevel)
 	{
 		return mSubresources[face * mNumMipLevels + mipLevel];
@@ -552,10 +564,8 @@ namespace bs
 									transferDstLayout, dstImage->getLayout(), dstRange);
 
 			// Notify the command buffer that these resources are being used on it
-			transferCB->getCB()->registerResource(srcImage, mAccessFlags, srcImage->getLayout(), srcRange, 
-				VulkanUseFlag::Read);
-			transferCB->getCB()->registerResource(dstImage, other->getAccessFlags(), dstImage->getLayout(), dstRange,
-				VulkanUseFlag::Write);
+			transferCB->getCB()->registerResource(srcImage, mAccessFlags, srcImage->getLayout(), VulkanUseFlag::Read);
+			transferCB->getCB()->registerResource(dstImage, other->getAccessFlags(), dstImage->getLayout(), VulkanUseFlag::Write);
 
 			// Need to wait if subresource we're reading from is being written, or if the subresource we're writing to is
 			// being accessed in any way
@@ -864,7 +874,7 @@ namespace bs
 
 				// Notify the command buffer that these resources are being used on it
 				transferCB->getCB()->registerResource(mStagingBuffer, VK_ACCESS_TRANSFER_READ_BIT, VulkanUseFlag::Read);
-				transferCB->getCB()->registerResource(image, mAccessFlags, image->getLayout(), range, VulkanUseFlag::Write);
+				transferCB->getCB()->registerResource(image, mAccessFlags, image->getLayout(), VulkanUseFlag::Write);
 
 				// We don't actually flush the transfer buffer here since it's an expensive operation, but it's instead
 				// done automatically before next "normal" command buffer submission.