1 files changed, 68 insertions, 147 deletions

diff --git a/src/video_core/shader_cache.h b/src/video_core/shader_cache.h
index 015a789d6..136fe294c 100644
--- a/src/video_core/shader_cache.h
+++ b/src/video_core/shader_cache.h
@@ -5,226 +5,147 @@
 #pragma once
 
 #include <algorithm>
+#include <array>
 #include <memory>
 #include <mutex>
+#include <span>
 #include <unordered_map>
 #include <utility>
 #include <vector>
 
-#include "common/assert.h"
 #include "common/common_types.h"
 #include "video_core/rasterizer_interface.h"
+#include "video_core/shader_environment.h"
+
+namespace Tegra {
+class MemoryManager;
+}
 
 namespace VideoCommon {
 
-template <class T>
+class GenericEnvironment;
+
+struct ShaderInfo {
+    u64 unique_hash{};
+    size_t size_bytes{};
+};
+
 class ShaderCache {
     static constexpr u64 PAGE_BITS = 14;
     static constexpr u64 PAGE_SIZE = u64(1) << PAGE_BITS;
 
+    static constexpr size_t NUM_PROGRAMS = 6;
+
     struct Entry {
         VAddr addr_start;
         VAddr addr_end;
-        T* data;
+        ShaderInfo* data;
 
         bool is_memory_marked = true;
 
-        constexpr bool Overlaps(VAddr start, VAddr end) const noexcept {
+        bool Overlaps(VAddr start, VAddr end) const noexcept {
             return start < addr_end && addr_start < end;
         }
     };
 
 public:
-    virtual ~ShaderCache() = default;
-
     /// @brief Removes shaders inside a given region
     /// @note Checks for ranges
     /// @param addr Start address of the invalidation
     /// @param size Number of bytes of the invalidation
-    void InvalidateRegion(VAddr addr, std::size_t size) {
-        std::scoped_lock lock{invalidation_mutex};
-        InvalidatePagesInRegion(addr, size);
-        RemovePendingShaders();
-    }
+    void InvalidateRegion(VAddr addr, size_t size);
 
     /// @brief Unmarks a memory region as cached and marks it for removal
     /// @param addr Start address of the CPU write operation
     /// @param size Number of bytes of the CPU write operation
-    void OnCPUWrite(VAddr addr, std::size_t size) {
-        std::lock_guard lock{invalidation_mutex};
-        InvalidatePagesInRegion(addr, size);
-    }
+    void OnCPUWrite(VAddr addr, size_t size);
 
     /// @brief Flushes delayed removal operations
-    void SyncGuestHost() {
-        std::scoped_lock lock{invalidation_mutex};
-        RemovePendingShaders();
-    }
+    void SyncGuestHost();
 
-    /// @brief Tries to obtain a cached shader starting in a given address
-    /// @note Doesn't check for ranges, the given address has to be the start of the shader
-    /// @param addr Start address of the shader, this doesn't cache for region
-    /// @return Pointer to a valid shader, nullptr when nothing is found
-    T* TryGet(VAddr addr) const {
-        std::scoped_lock lock{lookup_mutex};
+protected:
+    struct GraphicsEnvironments {
+        std::array<GraphicsEnvironment, NUM_PROGRAMS> envs;
+        std::array<Shader::Environment*, NUM_PROGRAMS> env_ptrs;
 
-        const auto it = lookup_cache.find(addr);
-        if (it == lookup_cache.end()) {
-            return nullptr;
+        std::span<Shader::Environment* const> Span() const noexcept {
+            return std::span(env_ptrs.begin(), std::ranges::find(env_ptrs, nullptr));
         }
-        return it->second->data;
-    }
-
-protected:
-    explicit ShaderCache(VideoCore::RasterizerInterface& rasterizer_) : rasterizer{rasterizer_} {}
+    };
 
-    /// @brief Register in the cache a given entry
-    /// @param data Shader to store in the cache
-    /// @param addr Start address of the shader that will be registered
-    /// @param size Size in bytes of the shader
-    void Register(std::unique_ptr<T> data, VAddr addr, std::size_t size) {
-        std::scoped_lock lock{invalidation_mutex, lookup_mutex};
+    explicit ShaderCache(VideoCore::RasterizerInterface& rasterizer_,
+                         Tegra::MemoryManager& gpu_memory_, Tegra::Engines::Maxwell3D& maxwell3d_,
+                         Tegra::Engines::KeplerCompute& kepler_compute_);
 
-        const VAddr addr_end = addr + size;
-        Entry* const entry = NewEntry(addr, addr_end, data.get());
+    /// @brief Update the hashes and information of shader stages
+    /// @param unique_hashes Shader hashes to store into when a stage is enabled
+    /// @return True no success, false on error
+    bool RefreshStages(std::array<u64, NUM_PROGRAMS>& unique_hashes);
 
-        const u64 page_end = (addr_end + PAGE_SIZE - 1) >> PAGE_BITS;
-        for (u64 page = addr >> PAGE_BITS; page < page_end; ++page) {
-            invalidation_cache[page].push_back(entry);
-        }
+    /// @brief Returns information about the current compute shader
+    /// @return Pointer to a valid shader, nullptr on error
+    const ShaderInfo* ComputeShader();
 
-        storage.push_back(std::move(data));
+    /// @brief Collect the current graphics environments
+    void GetGraphicsEnvironments(GraphicsEnvironments& result,
+                                 const std::array<u64, NUM_PROGRAMS>& unique_hashes);
 
-        rasterizer.UpdatePagesCachedCount(addr, size, 1);
-    }
+    Tegra::MemoryManager& gpu_memory;
+    Tegra::Engines::Maxwell3D& maxwell3d;
+    Tegra::Engines::KeplerCompute& kepler_compute;
 
-    /// @brief Called when a shader is going to be removed
-    /// @param shader Shader that will be removed
-    /// @pre invalidation_cache is locked
-    /// @pre lookup_mutex is locked
-    virtual void OnShaderRemoval([[maybe_unused]] T* shader) {}
+    std::array<const ShaderInfo*, NUM_PROGRAMS> shader_infos{};
+    bool last_shaders_valid = false;
 
 private:
+    /// @brief Tries to obtain a cached shader starting in a given address
+    /// @note Doesn't check for ranges, the given address has to be the start of the shader
+    /// @param addr Start address of the shader, this doesn't cache for region
+    /// @return Pointer to a valid shader, nullptr when nothing is found
+    ShaderInfo* TryGet(VAddr addr) const;
+
+    /// @brief Register in the cache a given entry
+    /// @param data Shader to store in the cache
+    /// @param addr Start address of the shader that will be registered
+    /// @param size Size in bytes of the shader
+    void Register(std::unique_ptr<ShaderInfo> data, VAddr addr, size_t size);
+
     /// @brief Invalidate pages in a given region
     /// @pre invalidation_mutex is locked
-    void InvalidatePagesInRegion(VAddr addr, std::size_t size) {
-        const VAddr addr_end = addr + size;
-        const u64 page_end = (addr_end + PAGE_SIZE - 1) >> PAGE_BITS;
-        for (u64 page = addr >> PAGE_BITS; page < page_end; ++page) {
-            auto it = invalidation_cache.find(page);
-            if (it == invalidation_cache.end()) {
-                continue;
-            }
-            InvalidatePageEntries(it->second, addr, addr_end);
-        }
-    }
+    void InvalidatePagesInRegion(VAddr addr, size_t size);
 
     /// @brief Remove shaders marked for deletion
     /// @pre invalidation_mutex is locked
-    void RemovePendingShaders() {
-        if (marked_for_removal.empty()) {
-            return;
-        }
-        // Remove duplicates
-        std::sort(marked_for_removal.begin(), marked_for_removal.end());
-        marked_for_removal.erase(std::unique(marked_for_removal.begin(), marked_for_removal.end()),
-                                 marked_for_removal.end());
-
-        std::vector<T*> removed_shaders;
-        removed_shaders.reserve(marked_for_removal.size());
-
-        std::scoped_lock lock{lookup_mutex};
-
-        for (Entry* const entry : marked_for_removal) {
-            removed_shaders.push_back(entry->data);
-
-            const auto it = lookup_cache.find(entry->addr_start);
-            ASSERT(it != lookup_cache.end());
-            lookup_cache.erase(it);
-        }
-        marked_for_removal.clear();
-
-        if (!removed_shaders.empty()) {
-            RemoveShadersFromStorage(std::move(removed_shaders));
-        }
-    }
+    void RemovePendingShaders();
 
     /// @brief Invalidates entries in a given range for the passed page
     /// @param entries         Vector of entries in the page, it will be modified on overlaps
     /// @param addr            Start address of the invalidation
     /// @param addr_end        Non-inclusive end address of the invalidation
     /// @pre invalidation_mutex is locked
-    void InvalidatePageEntries(std::vector<Entry*>& entries, VAddr addr, VAddr addr_end) {
-        std::size_t index = 0;
-        while (index < entries.size()) {
-            Entry* const entry = entries[index];
-            if (!entry->Overlaps(addr, addr_end)) {
-                ++index;
-                continue;
-            }
-
-            UnmarkMemory(entry);
-            RemoveEntryFromInvalidationCache(entry);
-            marked_for_removal.push_back(entry);
-        }
-    }
+    void InvalidatePageEntries(std::vector<Entry*>& entries, VAddr addr, VAddr addr_end);
 
     /// @brief Removes all references to an entry in the invalidation cache
     /// @param entry Entry to remove from the invalidation cache
     /// @pre invalidation_mutex is locked
-    void RemoveEntryFromInvalidationCache(const Entry* entry) {
-        const u64 page_end = (entry->addr_end + PAGE_SIZE - 1) >> PAGE_BITS;
-        for (u64 page = entry->addr_start >> PAGE_BITS; page < page_end; ++page) {
-            const auto entries_it = invalidation_cache.find(page);
-            ASSERT(entries_it != invalidation_cache.end());
-            std::vector<Entry*>& entries = entries_it->second;
-
-            const auto entry_it = std::find(entries.begin(), entries.end(), entry);
-            ASSERT(entry_it != entries.end());
-            entries.erase(entry_it);
-        }
-    }
+    void RemoveEntryFromInvalidationCache(const Entry* entry);
 
     /// @brief Unmarks an entry from the rasterizer cache
     /// @param entry Entry to unmark from memory
-    void UnmarkMemory(Entry* entry) {
-        if (!entry->is_memory_marked) {
-            return;
-        }
-        entry->is_memory_marked = false;
-
-        const VAddr addr = entry->addr_start;
-        const std::size_t size = entry->addr_end - addr;
-        rasterizer.UpdatePagesCachedCount(addr, size, -1);
-    }
+    void UnmarkMemory(Entry* entry);
 
     /// @brief Removes a vector of shaders from a list
     /// @param removed_shaders Shaders to be removed from the storage
     /// @pre invalidation_mutex is locked
     /// @pre lookup_mutex is locked
-    void RemoveShadersFromStorage(std::vector<T*> removed_shaders) {
-        // Notify removals
-        for (T* const shader : removed_shaders) {
-            OnShaderRemoval(shader);
-        }
-
-        // Remove them from the cache
-        const auto is_removed = [&removed_shaders](const std::unique_ptr<T>& shader) {
-            return std::find(removed_shaders.begin(), removed_shaders.end(), shader.get()) !=
-                   removed_shaders.end();
-        };
-        std::erase_if(storage, is_removed);
-    }
+    void RemoveShadersFromStorage(std::vector<ShaderInfo*> removed_shaders);
 
     /// @brief Creates a new entry in the lookup cache and returns its pointer
     /// @pre lookup_mutex is locked
-    Entry* NewEntry(VAddr addr, VAddr addr_end, T* data) {
-        auto entry = std::make_unique<Entry>(Entry{addr, addr_end, data});
-        Entry* const entry_pointer = entry.get();
+    Entry* NewEntry(VAddr addr, VAddr addr_end, ShaderInfo* data);
 
-        lookup_cache.emplace(addr, std::move(entry));
-        return entry_pointer;
-    }
+    /// @brief Create a new shader entry and register it
+    const ShaderInfo* MakeShaderInfo(GenericEnvironment& env, VAddr cpu_addr);
 
     VideoCore::RasterizerInterface& rasterizer;
 
@@ -233,7 +154,7 @@ private:
 
     std::unordered_map<u64, std::unique_ptr<Entry>> lookup_cache;
     std::unordered_map<u64, std::vector<Entry*>> invalidation_cache;
-    std::vector<std::unique_ptr<T>> storage;
+    std::vector<std::unique_ptr<ShaderInfo>> storage;
     std::vector<Entry*> marked_for_removal;
 };