Fixed false-positive image reuploads (#1557)

* Fixed false-positive image reuploads

* Fixed userfaultfd path, removed dead code, simplified calculations

* oopsie

* track potentially dirty images and hash them

* untrack only first page of the image in case of head access

* rebase, initialize hash, fix bounds check

* include image tail in the calculations

src/video_core/buffer_cache/buffer_cache.cpp

@@ -635,7 +635,7 @@ bool BufferCache::SynchronizeBufferFromImage(Buffer& buffer, VAddr device_addr,
                "Texel buffer aliases image subresources {:x} : {:x}", device_addr,
                image.info.guest_address);
     boost::container::small_vector<vk::BufferImageCopy, 8> copies;
-    u32 offset = buffer.Offset(image.cpu_addr);
+    u32 offset = buffer.Offset(image.info.guest_address);
     const u32 num_layers = image.info.resources.layers;
     const u32 max_offset = offset + size;
     for (u32 m = 0; m < image.info.resources.levels; m++) {

src/video_core/page_manager.cpp

@@ -114,8 +114,8 @@ struct PageManager::Impl {
 
             // Notify rasterizer about the fault.
             const VAddr addr = msg.arg.pagefault.address;
-            const VAddr addr_page = Common::AlignDown(addr, PAGESIZE);
-            rasterizer->InvalidateMemory(addr_page, PAGESIZE);
+            const VAddr addr_page = GetPageAddr(addr);
+            rasterizer->InvalidateMemory(addr, addr_page, PAGESIZE);
         }
     }
 
@@ -157,8 +157,8 @@ struct PageManager::Impl {
         const auto addr = reinterpret_cast<VAddr>(fault_address);
         const bool is_write = Common::IsWriteError(context);
         if (is_write && owned_ranges.find(addr) != owned_ranges.end()) {
-            const VAddr addr_aligned = Common::AlignDown(addr, PAGESIZE);
-            rasterizer->InvalidateMemory(addr_aligned, PAGESIZE);
+            const VAddr addr_aligned = GetPageAddr(addr);
+            rasterizer->InvalidateMemory(addr, addr_aligned, PAGESIZE);
             return true;
         }
         return false;
@@ -174,6 +174,14 @@ PageManager::PageManager(Vulkan::Rasterizer* rasterizer_)
 
 PageManager::~PageManager() = default;
 
+VAddr PageManager::GetPageAddr(VAddr addr) {
+    return Common::AlignDown(addr, PAGESIZE);
+}
+
+VAddr PageManager::GetNextPageAddr(VAddr addr) {
+    return Common::AlignUp(addr + 1, PAGESIZE);
+}
+
 void PageManager::OnGpuMap(VAddr address, size_t size) {
     impl->OnMap(address, size);
 }

src/video_core/page_manager.h

@@ -28,6 +28,9 @@ public:
     /// Increase/decrease the number of surface in pages touching the specified region
     void UpdatePagesCachedCount(VAddr addr, u64 size, s32 delta);
 
+    static VAddr GetPageAddr(VAddr addr);
+    static VAddr GetNextPageAddr(VAddr addr);
+
 private:
     struct Impl;
     std::unique_ptr<Impl> impl;

src/video_core/renderer_vulkan/vk_rasterizer.cpp

@@ -792,9 +792,9 @@ u32 Rasterizer::ReadDataFromGds(u32 gds_offset) {
     return value;
 }
 
-void Rasterizer::InvalidateMemory(VAddr addr, u64 size) {
-    buffer_cache.InvalidateMemory(addr, size);
-    texture_cache.InvalidateMemory(addr, size);
+void Rasterizer::InvalidateMemory(VAddr addr, VAddr addr_aligned, u64 size) {
+    buffer_cache.InvalidateMemory(addr_aligned, size);
+    texture_cache.InvalidateMemory(addr, addr_aligned, size);
 }
 
 void Rasterizer::MapMemory(VAddr addr, u64 size) {

src/video_core/renderer_vulkan/vk_rasterizer.h

@@ -46,7 +46,7 @@ public:
 
     void InlineData(VAddr address, const void* value, u32 num_bytes, bool is_gds);
     u32 ReadDataFromGds(u32 gsd_offset);
-    void InvalidateMemory(VAddr addr, u64 size);
+    void InvalidateMemory(VAddr addr, VAddr addr_aligned, u64 size);
     void MapMemory(VAddr addr, u64 size);
     void UnmapMemory(VAddr addr, u64 size);
 

src/video_core/texture_cache/image.cpp

@@ -144,8 +144,7 @@ void UniqueImage::Create(const vk::ImageCreateInfo& image_ci) {
 Image::Image(const Vulkan::Instance& instance_, Vulkan::Scheduler& scheduler_,
              const ImageInfo& info_)
     : instance{&instance_}, scheduler{&scheduler_}, info{info_},
-      image{instance->GetDevice(), instance->GetAllocator()}, cpu_addr{info.guest_address},
-      cpu_addr_end{cpu_addr + info.guest_size_bytes} {
+      image{instance->GetDevice(), instance->GetAllocator()} {
     mip_hashes.resize(info.resources.levels);
     ASSERT(info.pixel_format != vk::Format::eUndefined);
     // Here we force `eExtendedUsage` as don't know all image usage cases beforehand. In normal case

src/video_core/texture_cache/image.h

@@ -22,11 +22,12 @@ VK_DEFINE_HANDLE(VmaAllocator)
 namespace VideoCore {
 
 enum ImageFlagBits : u32 {
-    CpuDirty = 1 << 1, ///< Contents have been modified from the CPU
+    Empty = 0,
+    MaybeCpuDirty = 1 << 0, ///< The page this image is in was touched before the image address
+    CpuDirty = 1 << 1,      ///< Contents have been modified from the CPU
     GpuDirty = 1 << 2, ///< Contents have been modified from the GPU (valid data in buffer cache)
-    Dirty = CpuDirty | GpuDirty,
+    Dirty = MaybeCpuDirty | CpuDirty | GpuDirty,
     GpuModified = 1 << 3,    ///< Contents have been modified from the GPU
-    Tracked = 1 << 4,        ///< Writes and reads are being hooked from the CPU
     Registered = 1 << 6,     ///< True when the image is registered
     Picked = 1 << 7,         ///< Temporary flag to mark the image as picked
     MetaRegistered = 1 << 8, ///< True when metadata for this surface is known and registered
@@ -78,7 +79,9 @@ struct Image {
 
     [[nodiscard]] bool Overlaps(VAddr overlap_cpu_addr, size_t overlap_size) const noexcept {
         const VAddr overlap_end = overlap_cpu_addr + overlap_size;
-        return cpu_addr < overlap_end && overlap_cpu_addr < cpu_addr_end;
+        const auto image_addr = info.guest_address;
+        const auto image_end = info.guest_address + info.guest_size_bytes;
+        return image_addr < overlap_end && overlap_cpu_addr < image_end;
     }
 
     ImageViewId FindView(const ImageViewInfo& info) const {
@@ -99,14 +102,18 @@ struct Image {
     void CopyImage(const Image& image);
     void CopyMip(const Image& image, u32 mip);
 
+    bool IsTracked() {
+        return track_addr != 0 && track_addr_end != 0;
+    }
+
     const Vulkan::Instance* instance;
     Vulkan::Scheduler* scheduler;
     ImageInfo info;
     UniqueImage image;
     vk::ImageAspectFlags aspect_mask = vk::ImageAspectFlagBits::eColor;
     ImageFlagBits flags = ImageFlagBits::Dirty;
-    VAddr cpu_addr = 0;
-    VAddr cpu_addr_end = 0;
+    VAddr track_addr = 0;
+    VAddr track_addr_end = 0;
     std::vector<ImageViewInfo> image_view_infos;
     std::vector<ImageViewId> image_view_ids;
 
@@ -130,6 +137,7 @@ struct Image {
     std::vector<State> subresource_states{};
     boost::container::small_vector<u64, 14> mip_hashes{};
     u64 tick_accessed_last{0};
+    u64 hash{0};
 
     struct {
         union {

src/video_core/texture_cache/texture_cache.cpp

@@ -29,9 +29,12 @@ TextureCache::TextureCache(const Vulkan::Instance& instance_, Vulkan::Scheduler&
     info.UpdateSize();
     const ImageId null_id = slot_images.insert(instance, scheduler, info);
     ASSERT(null_id.index == NULL_IMAGE_ID.index);
-    const vk::Image& null_image = slot_images[null_id].image;
+    auto& img = slot_images[null_id];
+    const vk::Image& null_image = img.image;
     Vulkan::SetObjectName(instance.GetDevice(), null_image, "Null Image");
-    slot_images[null_id].flags = ImageFlagBits::Tracked;
+    img.flags = ImageFlagBits::Empty;
+    img.track_addr = img.info.guest_address;
+    img.track_addr_end = img.info.guest_address + img.info.guest_size_bytes;
 
     ImageViewInfo view_info;
     const auto null_view_id =
@@ -43,13 +46,43 @@ TextureCache::TextureCache(const Vulkan::Instance& instance_, Vulkan::Scheduler&
 
 TextureCache::~TextureCache() = default;
 
-void TextureCache::InvalidateMemory(VAddr address, size_t size) {
+void TextureCache::MarkAsMaybeDirty(ImageId image_id, Image& image) {
+    if (image.hash == 0) {
+        // Initialize hash
+        const u8* addr = std::bit_cast<u8*>(image.info.guest_address);
+        image.hash = XXH3_64bits(addr, image.info.guest_size_bytes);
+    }
+    image.flags |= ImageFlagBits::MaybeCpuDirty;
+    UntrackImage(image_id);
+}
+
+void TextureCache::InvalidateMemory(VAddr addr, VAddr page_addr, size_t size) {
     std::scoped_lock lock{mutex};
-    ForEachImageInRegion(address, size, [&](ImageId image_id, Image& image) {
-        // Ensure image is reuploaded when accessed again.
-        image.flags |= ImageFlagBits::CpuDirty;
-        // Untrack image, so the range is unprotected and the guest can write freely.
-        UntrackImage(image_id);
+    ForEachImageInRegion(page_addr, size, [&](ImageId image_id, Image& image) {
+        const auto image_begin = image.info.guest_address;
+        const auto image_end = image.info.guest_address + image.info.guest_size_bytes;
+        const auto page_end = page_addr + size;
+        if (image_begin <= addr && addr < image_end) {
+            // This image was definitely accessed by this page fault.
+            // Untrack image, so the range is unprotected and the guest can write freely
+            image.flags |= ImageFlagBits::CpuDirty;
+            UntrackImage(image_id);
+        } else if (page_end < image_end) {
+            // This page access may or may not modify the image.
+            // We should not mark it as dirty now. If it really was modified
+            // it will receive more invalidations on its other pages.
+            // Remove tracking from this page only.
+            UntrackImageHead(image_id);
+        } else if (image_begin < page_addr) {
+            // This page access does not modify the image but the page should be untracked.
+            // We should not mark this image as dirty now. If it really was modified
+            // it will receive more invalidations on its other pages.
+            UntrackImageTail(image_id);
+        } else {
+            // Image begins and ends on this page so it can not receive any more invalidations.
+            // We will check it's hash later to see if it really was modified.
+            MarkAsMaybeDirty(image_id, image);
+        }
     });
 }
 
@@ -415,6 +448,23 @@ void TextureCache::RefreshImage(Image& image, Vulkan::Scheduler* custom_schedule
         return;
     }
 
+    if (True(image.flags & ImageFlagBits::MaybeCpuDirty) &&
+        False(image.flags & ImageFlagBits::CpuDirty)) {
+        // The image size should be less than page size to be considered MaybeCpuDirty
+        // So this calculation should be very uncommon and reasonably fast
+        // For now we'll just check up to 64 first pixels
+        const auto addr = std::bit_cast<u8*>(image.info.guest_address);
+        const auto w = std::min(image.info.size.width, u32(8));
+        const auto h = std::min(image.info.size.height, u32(8));
+        const auto size = w * h * image.info.num_bits / 8;
+        const u64 hash = XXH3_64bits(addr, size);
+        if (image.hash == hash) {
+            image.flags &= ~ImageFlagBits::MaybeCpuDirty;
+            return;
+        }
+        image.hash = hash;
+    }
+
     const auto& num_layers = image.info.resources.layers;
     const auto& num_mips = image.info.resources.levels;
     ASSERT(num_mips == image.info.mips_layout.size());
@@ -425,14 +475,14 @@ void TextureCache::RefreshImage(Image& image, Vulkan::Scheduler* custom_schedule
         const u32 height = std::max(image.info.size.height >> m, 1u);
         const u32 depth =
             image.info.props.is_volume ? std::max(image.info.size.depth >> m, 1u) : 1u;
-        const auto& [mip_size, mip_pitch, mip_height, mip_ofs] = image.info.mips_layout[m];
+        const auto& mip = image.info.mips_layout[m];
 
         // Protect GPU modified resources from accidental CPU reuploads.
         const bool is_gpu_modified = True(image.flags & ImageFlagBits::GpuModified);
         const bool is_gpu_dirty = True(image.flags & ImageFlagBits::GpuDirty);
         if (is_gpu_modified && !is_gpu_dirty) {
             const u8* addr = std::bit_cast<u8*>(image.info.guest_address);
-            const u64 hash = XXH3_64bits(addr + mip_ofs, mip_size);
+            const u64 hash = XXH3_64bits(addr + mip.offset, mip.size);
             if (image.mip_hashes[m] == hash) {
                 continue;
             }
@@ -440,9 +490,9 @@ void TextureCache::RefreshImage(Image& image, Vulkan::Scheduler* custom_schedule
         }
 
         image_copy.push_back({
-            .bufferOffset = mip_ofs * num_layers,
-            .bufferRowLength = static_cast<u32>(mip_pitch),
-            .bufferImageHeight = static_cast<u32>(mip_height),
+            .bufferOffset = mip.offset * num_layers,
+            .bufferRowLength = static_cast<u32>(mip.pitch),
+            .bufferImageHeight = static_cast<u32>(mip.height),
             .imageSubresource{
                 .aspectMask = image.aspect_mask & ~vk::ImageAspectFlagBits::eStencil,
                 .mipLevel = m,
@@ -455,6 +505,7 @@ void TextureCache::RefreshImage(Image& image, Vulkan::Scheduler* custom_schedule
     }
 
     if (image_copy.empty()) {
+        image.flags &= ~ImageFlagBits::Dirty;
         return;
     }
 
@@ -500,7 +551,7 @@ void TextureCache::RegisterImage(ImageId image_id) {
     ASSERT_MSG(False(image.flags & ImageFlagBits::Registered),
                "Trying to register an already registered image");
     image.flags |= ImageFlagBits::Registered;
-    ForEachPage(image.cpu_addr, image.info.guest_size_bytes,
+    ForEachPage(image.info.guest_address, image.info.guest_size_bytes,
                 [this, image_id](u64 page) { page_table[page].push_back(image_id); });
 }
 
@@ -509,7 +560,7 @@ void TextureCache::UnregisterImage(ImageId image_id) {
     ASSERT_MSG(True(image.flags & ImageFlagBits::Registered),
                "Trying to unregister an already unregistered image");
     image.flags &= ~ImageFlagBits::Registered;
-    ForEachPage(image.cpu_addr, image.info.guest_size_bytes, [this, image_id](u64 page) {
+    ForEachPage(image.info.guest_address, image.info.guest_size_bytes, [this, image_id](u64 page) {
         const auto page_it = page_table.find(page);
         if (page_it == nullptr) {
             UNREACHABLE_MSG("Unregistering unregistered page=0x{:x}", page << PageShift);
@@ -527,25 +578,106 @@ void TextureCache::UnregisterImage(ImageId image_id) {
 
 void TextureCache::TrackImage(ImageId image_id) {
     auto& image = slot_images[image_id];
-    if (True(image.flags & ImageFlagBits::Tracked)) {
+    const auto image_begin = image.info.guest_address;
+    const auto image_end = image.info.guest_address + image.info.guest_size_bytes;
+    if (image_begin == image.track_addr && image_end == image.track_addr_end) {
         return;
     }
-    image.flags |= ImageFlagBits::Tracked;
-    tracker.UpdatePagesCachedCount(image.cpu_addr, image.info.guest_size_bytes, 1);
+
+    if (!image.IsTracked()) {
+        // Re-track the whole image
+        image.track_addr = image_begin;
+        image.track_addr_end = image_end;
+        tracker.UpdatePagesCachedCount(image_begin, image.info.guest_size_bytes, 1);
+    } else {
+        if (image_begin < image.track_addr) {
+            TrackImageHead(image_id);
+        }
+        if (image.track_addr_end < image_end) {
+            TrackImageTail(image_id);
+        }
+    }
+}
+
+void TextureCache::TrackImageHead(ImageId image_id) {
+    auto& image = slot_images[image_id];
+    const auto image_begin = image.info.guest_address;
+    if (image_begin == image.track_addr) {
+        return;
+    }
+    ASSERT(image.track_addr != 0 && image_begin < image.track_addr);
+    const auto size = image.track_addr - image_begin;
+    image.track_addr = image_begin;
+    tracker.UpdatePagesCachedCount(image_begin, size, 1);
+}
+
+void TextureCache::TrackImageTail(ImageId image_id) {
+    auto& image = slot_images[image_id];
+    const auto image_end = image.info.guest_address + image.info.guest_size_bytes;
+    if (image_end == image.track_addr_end) {
+        return;
+    }
+    ASSERT(image.track_addr_end != 0 && image.track_addr_end < image_end);
+    const auto addr = image.track_addr_end;
+    const auto size = image_end - image.track_addr_end;
+    image.track_addr_end = image_end;
+    tracker.UpdatePagesCachedCount(addr, size, 1);
 }
 
 void TextureCache::UntrackImage(ImageId image_id) {
     auto& image = slot_images[image_id];
-    if (False(image.flags & ImageFlagBits::Tracked)) {
+    if (!image.IsTracked()) {
         return;
     }
-    image.flags &= ~ImageFlagBits::Tracked;
-    tracker.UpdatePagesCachedCount(image.cpu_addr, image.info.guest_size_bytes, -1);
+    const auto addr = image.track_addr;
+    const auto size = image.track_addr_end - image.track_addr;
+    image.track_addr = 0;
+    image.track_addr_end = 0;
+    if (size != 0) {
+        tracker.UpdatePagesCachedCount(addr, size, -1);
+    }
+}
+
+void TextureCache::UntrackImageHead(ImageId image_id) {
+    auto& image = slot_images[image_id];
+    const auto image_begin = image.info.guest_address;
+    if (!image.IsTracked() || image_begin < image.track_addr) {
+        return;
+    }
+    const auto addr = tracker.GetNextPageAddr(image_begin);
+    const auto size = addr - image_begin;
+    image.track_addr = addr;
+    if (image.track_addr == image.track_addr_end) {
+        // This image spans only 2 pages and both are modified,
+        // but the image itself was not directly affected.
+        // Cehck its hash later.
+        MarkAsMaybeDirty(image_id, image);
+    }
+    tracker.UpdatePagesCachedCount(image_begin, size, -1);
+}
+
+void TextureCache::UntrackImageTail(ImageId image_id) {
+    auto& image = slot_images[image_id];
+    const auto image_end = image.info.guest_address + image.info.guest_size_bytes;
+    if (!image.IsTracked() || image.track_addr_end < image_end) {
+        return;
+    }
+    ASSERT(image.track_addr_end != 0);
+    const auto addr = tracker.GetPageAddr(image_end);
+    const auto size = image_end - addr;
+    image.track_addr_end = addr;
+    if (image.track_addr == image.track_addr_end) {
+        // This image spans only 2 pages and both are modified,
+        // but the image itself was not directly affected.
+        // Cehck its hash later.
+        MarkAsMaybeDirty(image_id, image);
+    }
+    tracker.UpdatePagesCachedCount(addr, size, -1);
 }
 
 void TextureCache::DeleteImage(ImageId image_id) {
     Image& image = slot_images[image_id];
-    ASSERT_MSG(False(image.flags & ImageFlagBits::Tracked), "Image was not untracked");
+    ASSERT_MSG(!image.IsTracked(), "Image was not untracked");
     ASSERT_MSG(False(image.flags & ImageFlagBits::Registered), "Image was not unregistered");
 
     // Remove any registered meta areas.

src/video_core/texture_cache/texture_cache.h

@@ -95,7 +95,7 @@ public:
     ~TextureCache();
 
     /// Invalidates any image in the logical page range.
-    void InvalidateMemory(VAddr address, size_t size);
+    void InvalidateMemory(VAddr addr, VAddr page_addr, size_t size);
 
     /// Marks an image as dirty if it exists at the provided address.
     void InvalidateMemoryFromGPU(VAddr address, size_t max_size);
@@ -242,9 +242,15 @@ private:
 
     /// Track CPU reads and writes for image
     void TrackImage(ImageId image_id);
+    void TrackImageHead(ImageId image_id);
+    void TrackImageTail(ImageId image_id);
 
     /// Stop tracking CPU reads and writes for image
     void UntrackImage(ImageId image_id);
+    void UntrackImageHead(ImageId image_id);
+    void UntrackImageTail(ImageId image_id);
+
+    void MarkAsMaybeDirty(ImageId image_id, Image& image);
 
     /// Removes the image and any views/surface metas that reference it.
     void DeleteImage(ImageId image_id);