Preserve storage images across render target transitions

src/graphics/host_gpu/renderer/imageInfo.h

@@ -173,7 +173,13 @@ struct VideoOutPixelFormatInfo {
 
 enum class DepthOverlap : uint8_t { None, RetireSampled, Unsupported };
 enum class DepthTransitionSource : uint8_t { None, Guest, Native };
-enum class RenderTargetOverlap : uint8_t { None, RetireSampled, RetireTarget, Unsupported };
+enum class RenderTargetOverlap : uint8_t {
+	None,
+	RetireSampled,
+	PreserveStorage,
+	RetireTarget,
+	Unsupported
+};
 enum class SampledOverlap : uint8_t { None, ReadOnlyAlias, Unsupported };
 enum class StorageSampledOverlap : uint8_t { None, ExactImage, Unsupported };
 enum class HostWriteOverlap : uint8_t { None, InvalidateImage, Unsupported };
@@ -474,6 +480,27 @@ ClassifyRenderTargetOverlap(const ImageInfo& sampled, bool sampled_gpu_modified,
 	           : RenderTargetOverlap::Unsupported;
 }
 
+[[nodiscard]] inline RenderTargetOverlap ClassifyStorageRenderTargetOverlap(
+    const ImageInfo& storage, VkFormat storage_format, bool storage_gpu_modified,
+    bool storage_buffer_modified, bool storage_cpu_dirty, bool same_context,
+    const RenderTargetInfo& target) {
+	if (!ImagePageRangesOverlap(storage.address, storage.size, target.address, target.size)) {
+		return RenderTargetOverlap::None;
+	}
+	const bool exact_native_image =
+	    storage.address == target.address && storage.size == target.size &&
+	    storage_format == target.format && storage.width == target.width &&
+	    storage.height == target.height && storage.pitch == target.pitch &&
+	    storage.base_level == 0 && storage.levels == 1 && storage.view_levels == 1 &&
+	    storage.tile == target.tile_mode && storage.depth == 1 &&
+	    storage.type == Prospero::GpuEnumValue(Prospero::ImageType::kColor2D) &&
+	    storage.base_array == 0 && target.levels == 1;
+	return exact_native_image && storage_gpu_modified && !storage_buffer_modified &&
+	               !storage_cpu_dirty && same_context
+	           ? RenderTargetOverlap::PreserveStorage
+	           : RenderTargetOverlap::Unsupported;
+}
+
 [[nodiscard]] inline bool IsRgba8SrgbViewFormat(VkFormat format) noexcept {
 	return format == VK_FORMAT_R8G8B8A8_UNORM || format == VK_FORMAT_R8G8B8A8_SRGB ||
 	       format == VK_FORMAT_B8G8R8A8_UNORM || format == VK_FORMAT_B8G8R8A8_SRGB;

src/graphics/host_gpu/renderer/textureCache.cpp

@@ -527,13 +527,15 @@ struct TextureCache::ReadbackWorker {
 	std::thread          thread;
 };
 
-void TextureCache::RetireImages(const std::vector<CachedImage*>& retire) {
+void TextureCache::RetireImages(const std::vector<CachedImage*>& retire,
+                                const CachedImage* native_storage_source) {
 	if (retire.empty()) {
 		return;
 	}
 	// FindTexture and FindRenderTarget retain cache records on their recording commands. Erasing
 	// the lookup owner therefore defers Vulkan destruction until the last referencing fence.
-	size_t removed = 0;
+	size_t removed                = 0;
+	bool   native_storage_retired = false;
 	for (auto it = m_images.begin(); it != m_images.end();) {
 		if (std::find(retire.begin(), retire.end(), it->get()) == retire.end()) {
 			++it;
@@ -541,7 +543,22 @@ void TextureCache::RetireImages(const std::vector<CachedImage*>& retire) {
 		}
 		const bool sampled = (*it)->kind == CachedImage::Kind::Texture;
 		const bool target  = (*it)->kind == CachedImage::Kind::RenderTarget;
-		if ((!sampled && !target) || (*it)->gpu_modified || (target && (*it)->buffer_modified)) {
+		const bool native_storage = it->get() == native_storage_source;
+		if (native_storage) {
+			if ((*it)->kind != CachedImage::Kind::StorageTexture || !(*it)->gpu_modified ||
+			    (*it)->buffer_modified || (*it)->info.IsCpuDirty() ||
+			    !m_memory_tracker.IsRegionGpuModified((*it)->Address(), (*it)->Size())) {
+				EXIT("TextureCache: invalid native storage-image retirement, kind=%u "
+				     "gpu_modified=%d buffer_modified=%d cpu_dirty=%d tracker_gpu_modified=%d\n",
+				     static_cast<uint32_t>((*it)->kind), (*it)->gpu_modified,
+				     (*it)->buffer_modified, (*it)->info.IsCpuDirty(),
+				     m_memory_tracker.IsRegionGpuModified((*it)->Address(), (*it)->Size()));
+			}
+			(*it)->gpu_modified  = false;
+			native_storage_retired = true;
+		}
+		if ((!sampled && !target && !native_storage) || (*it)->gpu_modified ||
+		    (target && (*it)->buffer_modified)) {
 			EXIT("TextureCache: invalid image retirement, kind=%u gpu_modified=%d "
 			     "buffer_modified=%d\n",
 			     static_cast<uint32_t>((*it)->kind), (*it)->gpu_modified, (*it)->buffer_modified);
@@ -553,7 +570,11 @@ void TextureCache::RetireImages(const std::vector<CachedImage*>& retire) {
 				     "addr=0x%016" PRIx64 " size=0x%016" PRIx64 "\n",
 				     (*it)->Address(range), (*it)->Size(range));
 			}
-			m_memory_tracker.UntrackMemory((*it)->Address(range), (*it)->Size(range));
+			// An exact storage-to-target transition transfers this registration to the new cache
+			// record. Keeping it GPU-owned also keeps guest page protection continuous.
+			if (!native_storage) {
+				m_memory_tracker.UntrackMemory((*it)->Address(range), (*it)->Size(range));
+			}
 		}
 		it = m_images.erase(it);
 		removed++;
@@ -562,6 +583,9 @@ void TextureCache::RetireImages(const std::vector<CachedImage*>& retire) {
 		EXIT("TextureCache: image retirement set mismatch, requested=%zu removed=%zu\n",
 		     retire.size(), removed);
 	}
+	if (native_storage_source != nullptr && !native_storage_retired) {
+		EXIT("TextureCache: native storage-image retirement source was not retired\n");
+	}
 }
 
 namespace {
@@ -1759,25 +1783,53 @@ RenderTextureVulkanImage* TextureCache::FindRenderTarget(CommandBuffer*
 		command->RetainResourceUntilFence(match);
 		return static_cast<RenderTextureVulkanImage*>(match->image);
 	}
-	std::vector<CachedImage*> retire;
+	std::vector<CachedImage*>       retire;
+	std::shared_ptr<CachedImage> native_storage_source;
 	for (const auto& entry: m_images) {
 		auto& cached = *entry;
 		if (!cached.OverlapsRange(info.address, info.size, true)) {
 			continue;
 		}
-		const auto overlap =
+		RenderTargetOverlap overlap = RenderTargetOverlap::Unsupported;
-		    cached.kind == CachedImage::Kind::Texture
+		switch (cached.kind) {
-		        ? ClassifyRenderTargetOverlap(cached.info, cached.gpu_modified, cached.ctx == ctx,
+			case CachedImage::Kind::Texture:
-		                                      info)
+				overlap = ClassifyRenderTargetOverlap(cached.info, cached.gpu_modified,
-		    : cached.kind == CachedImage::Kind::RenderTarget
+				                                      cached.ctx == ctx, info);
-		        ? ClassifyRenderTargetOverlap(cached.target, cached.gpu_modified,
+				break;
-		                                      cached.buffer_modified, cached.ctx == ctx, info)
+			case CachedImage::Kind::StorageTexture:
-		        : RenderTargetOverlap::Unsupported;
+				overlap = ClassifyStorageRenderTargetOverlap(
-		const bool retire_sampled = cached.kind == CachedImage::Kind::Texture &&
+				    cached.info, cached.image->format, cached.gpu_modified,
-		                            overlap == RenderTargetOverlap::RetireSampled;
+				    cached.buffer_modified,
-		const bool retire_target  = cached.kind == CachedImage::Kind::RenderTarget &&
+				    cached.info.IsCpuDirty() ||
-		                            overlap == RenderTargetOverlap::RetireTarget;
+				        m_memory_tracker.IsRegionCpuModified(cached.info.address, cached.info.size),
-		if (!retire_sampled && !retire_target) {
+				    cached.ctx == ctx, info);
+				break;
+			case CachedImage::Kind::RenderTarget:
+				overlap = ClassifyRenderTargetOverlap(cached.target, cached.gpu_modified,
+				                                      cached.buffer_modified, cached.ctx == ctx, info);
+				break;
+			case CachedImage::Kind::DepthTarget:
+			case CachedImage::Kind::VideoOut: break;
+		}
+		bool supported = false;
+		switch (overlap) {
+			case RenderTargetOverlap::RetireSampled:
+				supported = cached.kind == CachedImage::Kind::Texture;
+				break;
+			case RenderTargetOverlap::PreserveStorage:
+				supported = cached.kind == CachedImage::Kind::StorageTexture &&
+				            native_storage_source == nullptr;
+				if (supported) {
+					native_storage_source = entry;
+				}
+				break;
+			case RenderTargetOverlap::RetireTarget:
+				supported = cached.kind == CachedImage::Kind::RenderTarget;
+				break;
+			case RenderTargetOverlap::None:
+			case RenderTargetOverlap::Unsupported: break;
+		}
+		if (!supported) {
 			EXIT("TextureCache: unsupported render-target alias, requested=0x%016" PRIx64
 			     "+0x%016" PRIx64 " existing_kind=%u existing=0x%016" PRIx64 "+0x%016" PRIx64
 			     " gpu_modified=%d same_context=%d"
@@ -1809,18 +1861,37 @@ RenderTextureVulkanImage* TextureCache::FindRenderTarget(CommandBuffer*
 		     info.address, info.size, retired.address, retired.size, retained.address,
 		     retained.size);
 	}
-	RetireImages(retire);
+	RetireImages(retire, native_storage_source.get());
 	auto cached    = std::make_shared<CachedImage>();
 	cached->kind   = CachedImage::Kind::RenderTarget;
 	cached->target = info;
 	cached->ctx    = ctx;
 	cached->image  = CreateRenderTarget(ctx, info, &cached->memory);
-	m_memory_tracker.ForEachUploadRange(
+	const bool preserve_native = native_storage_source != nullptr;
-	    info.address, info.size, false, [](uint64_t, uint64_t) noexcept {},
+	if (preserve_native) {
-	    [&]() noexcept {
+		command->RetainResourceUntilFence(native_storage_source);
-		    UploadRenderTarget(ctx, static_cast<RenderTextureVulkanImage*>(cached->image), info,
+		ImageImageCopy region {};
-		                       false);
+		region.src_image = native_storage_source->image;
-	    });
+		region.src_level = 0;
+		region.dst_level = 0;
+		region.width     = info.width;
+		region.height    = info.height;
+		UtilImageToImage(command, std::span<const ImageImageCopy>(&region, 1), cached->image,
+		                 static_cast<uint64_t>(RENDER_COLOR_IMAGE_LAYOUT));
+	} else {
+		m_memory_tracker.ForEachUploadRange(
+		    info.address, info.size, false, [](uint64_t, uint64_t) noexcept {},
+		    [&]() noexcept {
+			    UploadRenderTarget(ctx, static_cast<RenderTextureVulkanImage*>(cached->image), info,
+			                       false);
+		    });
+	}
+	cached->gpu_modified = preserve_native;
+	if (preserve_native) {
+		LOGF("TextureCache: preserved storage image through native render-target copy, "
+		     "addr=0x%016" PRIx64 " size=0x%016" PRIx64 " format=%d extent=%ux%u\n",
+		     info.address, info.size, static_cast<int>(info.format), info.width, info.height);
+	}
 	auto* image = static_cast<RenderTextureVulkanImage*>(cached->image);
 	command->RetainResourceUntilFence(cached);
 	m_images.push_back(std::move(cached));

src/graphics/host_gpu/renderer/textureCache.h

@@ -115,7 +115,8 @@ private:
 	                                                              uint64_t size) const;
 	void                       RequireNoMetaOverlapLocked(uint64_t vaddr, uint64_t size) const;
 	void                       MarkSampledAliasesCpuDirtyLocked(uint64_t vaddr, uint64_t size);
-	void                       RetireImages(const std::vector<CachedImage*>& retire);
+	void RetireImages(const std::vector<CachedImage*>& retire,
+	                  const CachedImage* native_storage_source = nullptr);
 	void                       SynchronizeRenderTargetToBufferLocked(CachedImage& cached);
 
 	Common::Mutex                             m_dummy_mutex;

tests/ShaderRecompilerComputeTests.cpp

@@ -10539,6 +10539,52 @@ void CheckImageOverlapResolution() {
           ClassifyRenderTargetOverlap(sampled, true, true, target) ==
               RenderTargetOverlap::Unsupported,
           "GPU-owned sampled allocation was retired for a render target");
+
+  ImageInfo storage{};
+  storage.address = 0x112cd0000ull;
+  storage.size = 0xff0000;
+  storage.format = Prospero::GpuEnumValue(Prospero::BufferFormat::k16_16_16_16Float);
+  storage.width = 1920;
+  storage.height = 1080;
+  storage.pitch = 1920;
+  storage.tile = Prospero::GpuEnumValue(Prospero::TileMode::kRenderTarget);
+  storage.type = Prospero::GpuEnumValue(Prospero::ImageType::kColor2D);
+  RenderTargetInfo storage_target{};
+  storage_target.address = storage.address;
+  storage_target.size = storage.size;
+  storage_target.format = VK_FORMAT_R16G16B16A16_SFLOAT;
+  storage_target.width = storage.width;
+  storage_target.height = storage.height;
+  storage_target.pitch = storage.pitch;
+  storage_target.bytes_per_element = 8;
+  storage_target.tile_mode = storage.tile;
+  Require("ImageOverlapResolution", "storage render-target native transition",
+          ClassifyStorageRenderTargetOverlap(storage, storage_target.format, true, false, false,
+                                             true, storage_target) ==
+              RenderTargetOverlap::PreserveStorage,
+          "exact GPU-owned RGBA16F storage image was not preserved for a render target");
+  auto mismatched_storage_target = storage_target;
+  mismatched_storage_target.width--;
+  Require("ImageOverlapResolution", "storage render-target guards",
+          ClassifyStorageRenderTargetOverlap(storage, storage_target.format, false, false, false,
+                                             true, storage_target) ==
+                  RenderTargetOverlap::Unsupported &&
+              ClassifyStorageRenderTargetOverlap(storage, storage_target.format, true, true, false,
+                                                 true, storage_target) ==
+                  RenderTargetOverlap::Unsupported &&
+              ClassifyStorageRenderTargetOverlap(storage, storage_target.format, true, false, true,
+                                                 true, storage_target) ==
+                  RenderTargetOverlap::Unsupported &&
+              ClassifyStorageRenderTargetOverlap(storage, storage_target.format, true, false, false,
+                                                 false, storage_target) ==
+                  RenderTargetOverlap::Unsupported &&
+              ClassifyStorageRenderTargetOverlap(storage, VK_FORMAT_R32G32B32A32_SFLOAT, true,
+                                                 false, false, true, storage_target) ==
+                  RenderTargetOverlap::Unsupported &&
+              ClassifyStorageRenderTargetOverlap(storage, storage_target.format, true, false, false,
+                                                 true, mismatched_storage_target) ==
+                  RenderTargetOverlap::Unsupported,
+          "unsupported storage-to-render-target transition was admitted");
   Require("ImageOverlapResolution", "render target context",
           ClassifyRenderTargetOverlap(sampled, false, false, target) ==
               RenderTargetOverlap::Unsupported,