refactor(kernel): 添加一个显存的全局缓存，避免反复的 h2d 拷贝

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Jan 23, 2024 · 41036a0 · 41036a0
1 parent 20788a3
commit 41036a0
Showing 1 changed file with 26 additions and 3 deletions.
diff --git a/src/04kernel/src/graph.cc b/src/04kernel/src/graph.cc
@@ -1,5 +1,22 @@
 #include "kernel/graph.h"
 
+namespace refactor {
+    struct DataKey {
+        Arc<hardware::Device> dev;
+        Arc<kernel::Blob> blob;
+        bool operator==(const DataKey &) const = default;// since C++20
+    };
+}// namespace refactor
+
+template<>
+struct std::hash<refactor::DataKey> {
+    std::size_t operator()(refactor::DataKey const &s) const noexcept {
+        auto hd = std::hash<decltype(s.dev)>()(s.dev),
+             hb = std::hash<decltype(s.blob)>()(s.blob);
+        return hd ^ (hb << 1);
+    }
+};
+
 namespace refactor::kernel {
 
     Graph::Graph(graph_topo::GraphTopo topology,
@@ -31,13 +48,19 @@ namespace refactor::kernel {
             _internal.edges,
             32);
 
+        static std::unordered_map<DataKey, Arc<hardware::Device::Blob>> CACHE;
+
         for (auto i : range0_(edges_.size())) {
             auto const &edge = _internal.edges[i];
             edges_[i].name = edge.name;
             if (edge.data) {
-                auto blob = device->malloc(edge.size);
-                blob->copyFromHost(edge.data->get<void>());
-                edges_[i].blob = std::move(blob);
+                auto it = CACHE.find({device, edge.data});
+                if (it == CACHE.end()) {
+                    auto blob = device->malloc(edge.size);
+                    blob->copyFromHost(edge.data->get<void>());
+                    std::tie(it, std::ignore) = CACHE.emplace(DataKey{device, edge.data}, std::move(blob));
+                }
+                edges_[i].blob = it->second;
             } else if (edges_[i].stackOffset == SIZE_MAX - 1) {
                 edges_[i].blob = device->malloc(edge.size);
             }