feat(libSceAgcDriver): storage image access to depth surfaces

core/libs/prx/libSceAgcDriver/Graphics/include/DepthSurface.hpp

@@ -12,6 +12,7 @@
 namespace AgcDriver::Graphics {
 
 class Texture;
+class StorageTexture;
 
 VkImageView DepthSurfaceView(const Context& context, const DepthTarget& target);
 bool SampleLocationsCompatibleDepth(const Context& context, VkFormat format, std::uint32_t samples);
@@ -36,6 +37,8 @@ bool HtileFillCovers(std::uint64_t htile, VkExtent2D extent, std::uint64_t addre
 std::shared_ptr<Texture> DepthSurfaceTexture(const Context& context, std::span<const std::uint32_t> words, const GuestTextureResource& resource, VkComponentMapping components);
 VkFormat HostDepthFormat(const Context& context, VkFormat format);
 float HostDepthBiasConstant(const Context& context, VkFormat format, float constant);
+std::shared_ptr<StorageTexture> DepthSurfaceStorage(const Context& context, std::span<const std::uint32_t> words, const GuestTextureResource& resource, std::uint32_t mip);
+bool DepthSurfaceStoreBack(const Context& context, const StorageTexture& storage);
 
 }
 

core/libs/prx/libSceAgcDriver/Graphics/src/DepthSurface.cpp

@@ -153,6 +153,56 @@ class DepthSurface {
     float clearDepth = 0.0f;
     std::uint8_t clearStencil = 0;
 
+    std::shared_ptr<StorageTexture> Storage(std::span<const std::uint32_t> words, const GuestTextureResource& resource, std::uint32_t mip) {
+        if (resource.width != target.extent.width || resource.height != target.extent.height) return nullptr;
+        const bool stencil = target.stencilAddress != 0 && resource.baseAddress == target.stencilAddress;
+        const bool d16 = target.format == VK_FORMAT_D16_UNORM || target.format == VK_FORMAT_D16_UNORM_S8_UINT;
+        const auto expected = stencil ? VK_FORMAT_R8_UINT : d16 ? VK_FORMAT_R16_UNORM : VK_FORMAT_R32_SFLOAT;
+        const bool depthBits = !stencil && words.size() >= 4 && ShaderRecompiler::DepthBitsTextureWidth(words[1], words[3]) == (d16 ? 16u : 32u);
+        if (ResolveTextureFormat(resource.format) != expected && !depthBits) return nullptr;
+        if (resource.dimension != TextureDimension::k2D) {
+            char text[256];
+            std::snprintf(text, sizeof(text), "AGC graphics: storage image access to depth/stencil surface 0x%llx is not implemented (storage=%ux%u guest format=%u mip=%u dimension=%d tile=%d)", static_cast<unsigned long long>(resource.baseAddress), resource.width, resource.height, resource.format, mip, static_cast<int>(resource.dimension), static_cast<int>(resource.tileMode));
+            throw std::runtime_error(text);
+        }
+        Require(target.samples == 1, "storage image access to a multisampled depth surface is not implemented");
+        if (!stencil && hostFormat != target.format) {
+            char text[224];
+            std::snprintf(text, sizeof(text), "AGC graphics: storage image access to the depth plane of depth surface 0x%llx (vk format %d) stored as vk format %d is not implemented", static_cast<unsigned long long>(target.address), static_cast<int>(target.format), static_cast<int>(hostFormat));
+            throw std::runtime_error(text);
+        }
+        if (mip != 0 || resource.baseLevel != 0) {
+            char text[224];
+            std::snprintf(text, sizeof(text), "AGC graphics: storage image access to level %u of the %s plane of depth surface 0x%llx is not implemented", mip, stencil ? "stencil" : "depth", static_cast<unsigned long long>(target.address));
+            throw std::runtime_error(text);
+        }
+        ApplyFastClear();
+        auto& plane = stencil ? stencilPlane : depthPlane;
+        if (plane.storage == nullptr || plane.storage->Descriptor().format != resource.format || plane.storage->Descriptor().tileMode != resource.tileMode || plane.storage->Descriptor().baseAddress != resource.baseAddress) {
+            std::erase_if(plane.replaced, [](const auto& old) { return old.expired(); });
+            if (plane.storage != nullptr) plane.replaced.push_back(plane.storage);
+            plane.storage = StorageTexture::Create(context, *context.detiler, resource, 0);
+        }
+        copyPlane(plane, stencil, true);
+        return plane.storage;
+    }
+
+    bool StoreBack(const StorageTexture& candidate) {
+        for (const bool stencil : {false, true}) {
+            auto& plane = stencil ? stencilPlane : depthPlane;
+            if (plane.storage.get() == &candidate) {
+                copyPlane(plane, stencil, false);
+                written = ++depthWrites;
+                return true;
+            }
+            if (std::none_of(plane.replaced.begin(), plane.replaced.end(), [&](const auto& old) { return old.lock().get() == &candidate; })) continue;
+            char text[288];
+            std::snprintf(text, sizeof(text), "AGC graphics: storage image writes to the %s plane of depth surface 0x%llx through a view that another view of the plane (guest format %u, tile mode %d) replaced are not implemented", stencil ? "stencil" : "depth", static_cast<unsigned long long>(target.address), plane.storage->Descriptor().format, static_cast<int>(plane.storage->Descriptor().tileMode));
+            throw std::runtime_error(text);
+        }
+        return false;
+    }
+
     const Context context;
     const DepthTarget target;
     const VkFormat hostFormat;
@@ -161,8 +211,41 @@ class DepthSurface {
     VkDeviceMemory memory = VK_NULL_HANDLE;
 
 private:
+    struct PlaneStorage {
+        std::shared_ptr<StorageTexture> storage;
+        std::vector<std::weak_ptr<StorageTexture>> replaced;
+        std::unique_ptr<DeviceBuffer> staging;
+    };
     std::map<std::array<std::uint32_t, 12>, std::shared_ptr<Texture>> textures;
+    PlaneStorage depthPlane;
+    PlaneStorage stencilPlane;
+
+    void copyPlane(PlaneStorage& plane, bool stencil, bool toStorage) {
+        auto* recorder = Recorder::Active();
+        Require(recorder != nullptr, "depth surface storage copies need an active recorder");
+        const bool d16 = target.format == VK_FORMAT_D16_UNORM || target.format == VK_FORMAT_D16_UNORM_S8_UINT;
+        const auto bytes = static_cast<std::size_t>(target.extent.width) * target.extent.height * (stencil ? 1u : d16 ? 2u : 4u);
+        if (plane.staging == nullptr) plane.staging = std::make_unique<DeviceBuffer>(context, bytes, VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);
+        const auto commands = recorder->Commands();
+        recorder->Keep(plane.storage);
+        const auto timing = recorder->BeginGpuTiming(Recorder::CommandClass::Copy);
+        VkBufferImageCopy planeRegion{0, 0, 0, {stencil ? VK_IMAGE_ASPECT_STENCIL_BIT : VK_IMAGE_ASPECT_DEPTH_BIT, 0, 0, 1}, {0, 0, 0}, {target.extent.width, target.extent.height, 1}};
+        VkBufferImageCopy colorRegion = planeRegion;
+        colorRegion.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
+        const auto sourceImage = toStorage ? image : plane.storage->Image();
+        const auto destinationImage = toStorage ? plane.storage->Image() : image;
+        RecordMemoryBarrier(context, commands, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_MEMORY_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT);
+        context.Function<PFN_vkCmdCopyImageToBuffer>("vkCmdCopyImageToBuffer")(commands, sourceImage, VK_IMAGE_LAYOUT_GENERAL, plane.staging->Handle(), 1, toStorage ? &planeRegion : &colorRegion);
+        RecordMemoryBarrier(context, commands, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT | VK_ACCESS_TRANSFER_WRITE_BIT);
+        context.Function<PFN_vkCmdCopyBufferToImage>("vkCmdCopyBufferToImage")(commands, plane.staging->Handle(), destinationImage, VK_IMAGE_LAYOUT_GENERAL, 1, toStorage ? &colorRegion : &planeRegion);
+        RecordMemoryBarrier(context, commands, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT);
+        Recorder::CountBarriers(Recorder::CommandClass::Copy, 3);
+        recorder->EndGpuTiming(timing, bytes);
+    }
+
     void release() noexcept {
+        depthPlane = {};
+        stencilPlane = {};
         textures.clear();
         if (view) context.Function<PFN_vkDestroyImageView>("vkDestroyImageView")(context.device, view, nullptr);
         if (image) context.Function<PFN_vkDestroyImage>("vkDestroyImage")(context.device, image, nullptr);
@@ -598,6 +681,23 @@ bool DepthSurfaceAt(const Context& context, std::span<const std::uint32_t> words
     return ResolveTextureFormat(resource.format) == expected || depthBits;
 }
 
+std::shared_ptr<StorageTexture> DepthSurfaceStorage(const Context& context, std::span<const std::uint32_t> words, const GuestTextureResource& resource, std::uint32_t mip) {
+    std::lock_guard lock(surfacesMutex());
+    const auto& list = surfaces();
+    const auto found = std::find_if(list.rbegin(), list.rend(), [&](const auto& surface) {
+        return surface->context.device == context.device && (surface->target.address == resource.baseAddress || (surface->target.stencilAddress != 0 && surface->target.stencilAddress == resource.baseAddress));
+    });
+    return found == list.rend() ? nullptr : (*found)->Storage(words, resource, mip);
+}
+
+bool DepthSurfaceStoreBack(const Context& context, const StorageTexture& storage) {
+    std::lock_guard lock(surfacesMutex());
+    for (const auto& surface : surfaces()) {
+        if (surface->context.device == context.device && surface->StoreBack(storage)) return true;
+    }
+    return false;
+}
+
 bool DepthSurfaceAt(std::uint64_t address) {
     std::lock_guard lock(surfacesMutex());
     return std::any_of(surfaces().begin(), surfaces().end(), [&](const auto& surface) { return surface->target.address == address || surface->target.stencilAddress == address; });

core/libs/prx/libSceAgcDriver/Graphics/src/ShaderResources.cpp

@@ -837,11 +837,7 @@ GuestTextureResource StorageSurface(const Context& context, const GuestTextureRe
 // `guestBytes` is the surface size when the caller described the surface already (0: described here).
 std::shared_ptr<StorageTexture> cachedStorageTexture(const Context& context, std::span<const std::uint32_t> words, const GuestTextureResource& viewed, std::uint32_t mip, std::uint64_t guestBytes) {
     RequireNoMultisampleTarget(viewed.baseAddress, guestBytes != 0 ? static_cast<std::size_t>(guestBytes) : 1u, "a storage image access");
-    if (DepthSurfaceAt(context, words, viewed)) {
-        char text[256];
-        std::snprintf(text, sizeof(text), "AGC graphics: storage image access to depth/stencil surface 0x%llx is not implemented (storage=%ux%u guest format=%u mip=%u dimension=%d tile=%d)", static_cast<unsigned long long>(viewed.baseAddress), viewed.width, viewed.height, viewed.format, mip, static_cast<int>(viewed.dimension), static_cast<int>(viewed.tileMode));
-        throw std::runtime_error(text);
-    }
+    if (auto depth = DepthSurfaceStorage(context, words, viewed, mip)) return depth;
     static const bool disabled = std::getenv("APS5_NO_TEXTURE_CACHE") != nullptr;
     if (disabled) return StorageTexture::Create(context, *context.detiler, viewed, mip);
     static_cast<void>(words);
@@ -3613,7 +3609,9 @@ void ShaderResources::WriteBack() {
     WriteBackBuffers();
     if (SkipWriteBack()) return;
     for (std::size_t index = 0; index < storageTextures.size(); ++index) {
-        if (storageWritten[index]) storageTextures[index]->MarkDirty();
+        if (!storageWritten[index]) continue;
+        storageTextures[index]->MarkDirty();
+        static_cast<void>(DepthSurfaceStoreBack(context, *storageTextures[index]));
     }
 }
 
@@ -3624,7 +3622,9 @@ void ShaderResources::MarkGpuWrites(Recorder& recorder) {
     recorder.NotePendingReads(guestMemory.InPlaceReads(), guestMemory.HoldsLease() ? Recorder::ReadKind::AddressBased : Recorder::ReadKind::DispatchElement);
     if (SkipWriteBack()) return;
     for (std::size_t index = 0; index < storageTextures.size(); ++index) {
-        if (storageWritten[index]) storageTextures[index]->MarkDirty();
+        if (!storageWritten[index]) continue;
+        storageTextures[index]->MarkDirty();
+        static_cast<void>(DepthSurfaceStoreBack(context, *storageTextures[index]));
     }
     // Written sub-ranges of buffers the GPU copied out of a host import go back into it by the GPU,
     // recorded here after the work: those regions then need no CPU write-back (HasCopiedWrites),

core/libs/prx/libSceAgcDriver/tests/DepthSurfaceReuse.cpp

@@ -7,6 +7,7 @@
 #include <map>
 #include <memory>
 #include <stdexcept>
+#include <string>
 #include <string_view>
 
 namespace {
@@ -192,6 +193,14 @@ void TestD16StencilFallback(const Context& context) {
     };
     Require(refused({d16s8, FormatR16Unorm, 256, 256, TileDepth64KB, Type2DArray, 1}), "two slices of a D16 surface stored as D24 must be refused, as the plane copy would read the wider image as 16-bit texels");
     Require(refused({d16s8, FormatR16Uint, 256, 256, TileStandard64KB}), "an R16_UINT view of a D16 surface stored as D24 must be refused");
+    std::string storageError;
+    try {
+        const auto words = tsharp({d16s8, FormatR16Unorm, 256, 256});
+        DepthSurfaceStorage(without, words, DecodeTextureResource(words), 0);
+    } catch (const std::runtime_error& error) {
+        storageError = error.what();
+    }
+    Require(storageError.find("depth plane") != std::string::npos && storageError.find("stored as") != std::string::npos, "storage access to the depth plane of a D16 surface stored as D24 must throw");
     ClearDepthSurfaces(without.device);
 
     static int onlyTag = 0;