3 files changed, 49 insertions, 10 deletions

diff --git a/src/video_core/control/channel_state_cache.cpp b/src/video_core/control/channel_state_cache.cpp
index f72a97b2f..ec7ba907c 100644
--- a/src/video_core/control/channel_state_cache.cpp
+++ b/src/video_core/control/channel_state_cache.cpp
@@ -1,5 +1,11 @@
 #include "video_core/control/channel_state_cache.inc"
 
 namespace VideoCommon {
+
+ChannelInfo::ChannelInfo(Tegra::Control::ChannelState& channel_state)
+    : maxwell3d{*channel_state.maxwell_3d}, kepler_compute{*channel_state.kepler_compute},
+      gpu_memory{*channel_state.memory_manager} {}
+
 template class VideoCommon::ChannelSetupCaches<VideoCommon::ChannelInfo>;
-}
+
+} // namespace VideoCommon
diff --git a/src/video_core/control/channel_state_cache.h b/src/video_core/control/channel_state_cache.h
index c8298c003..c51040c83 100644
--- a/src/video_core/control/channel_state_cache.h
+++ b/src/video_core/control/channel_state_cache.h
@@ -2,6 +2,7 @@
 
 #include <deque>
 #include <limits>
+#include <mutex>
 #include <unordered_map>
 
 #include "common/common_types.h"
@@ -41,9 +42,10 @@ template <class P>
 class ChannelSetupCaches {
 public:
     /// Operations for seting the channel of execution.
+    virtual ~ChannelSetupCaches();
 
     /// Create channel state.
-    void CreateChannel(Tegra::Control::ChannelState& channel);
+    virtual void CreateChannel(Tegra::Control::ChannelState& channel);
 
     /// Bind a channel for execution.
     void BindToChannel(s32 id);
@@ -51,18 +53,34 @@ public:
     /// Erase channel's state.
     void EraseChannel(s32 id);
 
+    Tegra::MemoryManager* GetFromID(size_t id) const {
+        std::unique_lock<std::mutex> lk(config_mutex);
+        const auto ref = address_spaces.find(id);
+        return ref->second.gpu_memory;
+    }
+
 protected:
     static constexpr size_t UNSET_CHANNEL{std::numeric_limits<size_t>::max()};
 
-    std::deque<P> channel_storage;
-    std::deque<size_t> free_channel_ids;
-    std::unordered_map<s32, size_t> channel_map;
-
     P* channel_state;
     size_t current_channel_id{UNSET_CHANNEL};
+    size_t current_address_space{};
     Tegra::Engines::Maxwell3D* maxwell3d;
     Tegra::Engines::KeplerCompute* kepler_compute;
     Tegra::MemoryManager* gpu_memory;
+
+    std::deque<P> channel_storage;
+    std::deque<size_t> free_channel_ids;
+    std::unordered_map<s32, size_t> channel_map;
+    struct AddresSpaceRef {
+        size_t ref_count;
+        size_t storage_id;
+        Tegra::MemoryManager* gpu_memory;
+    };
+    std::unordered_map<size_t, AddresSpaceRef> address_spaces;
+    mutable std::mutex config_mutex;
+
+    virtual void OnGPUASRegister([[maybe_unused]] size_t map_id) {}
 };
 
 } // namespace VideoCommon
diff --git a/src/video_core/control/channel_state_cache.inc b/src/video_core/control/channel_state_cache.inc
index 3eb73af9f..185eabc35 100644
--- a/src/video_core/control/channel_state_cache.inc
+++ b/src/video_core/control/channel_state_cache.inc
@@ -6,18 +6,18 @@
 
 namespace VideoCommon {
 
-ChannelInfo::ChannelInfo(Tegra::Control::ChannelState& channel_state)
-    : maxwell3d{*channel_state.maxwell_3d}, kepler_compute{*channel_state.kepler_compute},
-      gpu_memory{*channel_state.memory_manager} {}
+template <class P>
+ChannelSetupCaches<P>::~ChannelSetupCaches() = default;
 
 template <class P>
 void ChannelSetupCaches<P>::CreateChannel(struct Tegra::Control::ChannelState& channel) {
+    std::unique_lock<std::mutex> lk(config_mutex);
     ASSERT(channel_map.find(channel.bind_id) == channel_map.end() && channel.bind_id >= 0);
     auto new_id = [this, &channel]() {
         if (!free_channel_ids.empty()) {
             auto id = free_channel_ids.front();
             free_channel_ids.pop_front();
-            new (&channel_storage[id]) ChannelInfo(channel);
+            new (&channel_storage[id]) P(channel);
             return id;
         }
         channel_storage.emplace_back(channel);
@@ -27,11 +27,24 @@ void ChannelSetupCaches<P>::CreateChannel(struct Tegra::Control::ChannelState& c
     if (current_channel_id != UNSET_CHANNEL) {
         channel_state = &channel_storage[current_channel_id];
     }
+    auto as_it = address_spaces.find(channel.memory_manager->GetID());
+    if (as_it != address_spaces.end()) {
+      as_it->second.ref_count++;
+      return;
+    }
+    AddresSpaceRef new_gpu_mem_ref{
+      .ref_count = 1,
+      .storage_id = address_spaces.size(),
+      .gpu_memory = channel.memory_manager.get(),
+    };
+    address_spaces.emplace(channel.memory_manager->GetID(), new_gpu_mem_ref);
+    OnGPUASRegister(channel.memory_manager->GetID());
 }
 
 /// Bind a channel for execution.
 template <class P>
 void ChannelSetupCaches<P>::BindToChannel(s32 id) {
+    std::unique_lock<std::mutex> lk(config_mutex);
     auto it = channel_map.find(id);
     ASSERT(it != channel_map.end() && id >= 0);
     current_channel_id = it->second;
@@ -39,11 +52,13 @@ void ChannelSetupCaches<P>::BindToChannel(s32 id) {
     maxwell3d = &channel_state->maxwell3d;
     kepler_compute = &channel_state->kepler_compute;
     gpu_memory = &channel_state->gpu_memory;
+    current_address_space = gpu_memory->GetID();
 }
 
 /// Erase channel's channel_state.
 template <class P>
 void ChannelSetupCaches<P>::EraseChannel(s32 id) {
+    std::unique_lock<std::mutex> lk(config_mutex);
     const auto it = channel_map.find(id);
     ASSERT(it != channel_map.end() && id >= 0);
     const auto this_id = it->second;