perf: cache CPU position embeddings across DiT runners (#2211)

src/core/pe_cache.cpp

@@ -0,0 +1,30 @@
+#include "core/pe_cache.h"
+
+namespace sd {
+    PECacheKey& PECacheKey::add_shape(const ggml_tensor* tensor) {
+        add(tensor != nullptr);
+        if (tensor != nullptr) {
+            for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+                add(tensor->ne[i]);
+            }
+        }
+        return *this;
+    }
+
+    PECacheKey& PECacheKey::add_shapes(const std::vector<ggml_tensor*>& tensors) {
+        add(tensors.size());
+        for (const auto* tensor : tensors) {
+            add_shape(tensor);
+        }
+        return *this;
+    }
+
+    const Tensor<float>& PECache::get_or_create(const PECacheKey& key, const std::function<Tensor<float>()>& build) {
+        auto entry = entries_.find(key.name());
+        if (entry != entries_.end() && entry->second.key == key) {
+            return entry->second.value;
+        }
+        auto value = build();
+        return entries_.insert_or_assign(key.name(), Entry{key, std::move(value)}).first->second.value;
+    }
+}

src/core/pe_cache.h

@@ -0,0 +1,82 @@
+#ifndef __SD_CORE_PE_CACHE_H__
+#define __SD_CORE_PE_CACHE_H__
+
+#include <functional>
+#include <map>
+#include <set>
+#include <type_traits>
+
+#include "core/tensor.hpp"
+#include "ggml.h"
+
+namespace sd {
+    class PECacheKey {
+        std::string name_;
+        std::string fields_;
+
+    public:
+        explicit PECacheKey(std::string name)
+            : name_(std::move(name)) {}
+        const std::string& name() const { return name_; }
+
+        template <typename T, std::enable_if_t<std::is_arithmetic_v<T> || std::is_enum_v<T>, int> = 0>
+        PECacheKey& add(T value) {
+            if constexpr (std::is_floating_point_v<T>) {
+                const double field = static_cast<double>(value);
+                fields_.push_back('f');
+                fields_.append(reinterpret_cast<const char*>(&field), sizeof(field));
+            } else {
+                const uint64_t field = static_cast<uint64_t>(value);
+                fields_.push_back('i');
+                fields_.append(reinterpret_cast<const char*>(&field), sizeof(field));
+            }
+            return *this;
+        }
+
+        template <typename T>
+        PECacheKey& add(const std::vector<T>& values) {
+            fields_.push_back('v');
+            add(values.size());
+            for (const auto& value : values) {
+                add(value);
+            }
+            return *this;
+        }
+
+        template <typename T>
+        PECacheKey& add(const std::set<T>& values) {
+            return add(std::vector<T>(values.begin(), values.end()));
+        }
+
+        template <typename T>
+        PECacheKey& add(const Tensor<T>& values) {
+            add(values.shape());
+            for (int64_t i = 0; i < values.numel(); ++i) {
+                add(values.data()[i]);
+            }
+            return *this;
+        }
+
+        PECacheKey& add_shape(const ggml_tensor* tensor);
+        PECacheKey& add_shapes(const std::vector<ggml_tensor*>& tensors);
+
+        bool operator==(const PECacheKey& other) const {
+            return name_ == other.name_ && fields_ == other.fields_;
+        }
+    };
+
+    class PECache {
+        struct Entry {
+            PECacheKey key;
+            Tensor<float> value;
+        };
+
+        std::map<std::string, Entry> entries_;
+
+    public:
+        const Tensor<float>& get_or_create(const PECacheKey& key, const std::function<Tensor<float>()>& build);
+    };
+
+}
+
+#endif  // __SD_CORE_PE_CACHE_H__