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

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__