refactor: move model-specific operators and fusions into backend extensions (#2253)

src/core/ggml_compat.cpp

@@ -0,0 +1,185 @@
+#include "core/ggml_compat.h"
+#include "core/ggml_extend_backend.h"
+
+#ifdef SD_GGML_GRAPH_COMPAT
+// Upstream ggml still requires its private graph layout for these operations.
+#include "ggml-impl.h"
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+int ggml_graph_n_leafs(const ggml_cgraph* graph) {
+    return graph->n_leafs;
+}
+
+ggml_tensor* ggml_graph_leaf(const ggml_cgraph* graph, int index) {
+    if (index < 0) {
+        index += graph->n_leafs;
+    }
+    GGML_ASSERT(index >= 0 && index < graph->n_leafs);
+    return graph->leafs[index];
+}
+
+void ggml_graph_add_leaf(ggml_cgraph* graph, ggml_tensor* tensor) {
+    GGML_ASSERT(tensor && graph->n_leafs < graph->size);
+    graph->leafs[graph->n_leafs++] = tensor;
+}
+
+int32_t ggml_graph_use_count(const ggml_cgraph* graph, const ggml_tensor* tensor) {
+    if (!graph || !tensor || !graph->visited_hash_set.size) {
+        return 0;
+    }
+    const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor);
+    return index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index)
+               ? graph->use_counts[index]
+               : 0;
+}
+
+bool ggml_graph_cpy_mapped(const ggml_cgraph* src, ggml_cgraph* dst, ggml_graph_tensor_map_fn map, void* user_data) {
+    if (!src || !dst || src == dst || !map || dst->grads || dst->grad_accs ||
+        src->n_nodes > dst->size || src->n_leafs > dst->size) {
+        return false;
+    }
+    std::unordered_map<const ggml_tensor*, ggml_tensor*> mapping;
+    std::unordered_set<ggml_tensor*> mapped;
+    std::vector<const ggml_tensor*> sources;
+    bool valid = true;
+    auto remap = [&](const ggml_tensor* source) {
+        const auto found = mapping.find(source);
+        if (found != mapping.end()) {
+            return found->second;
+        }
+        auto* tensor = map(source, user_data);
+        if (tensor && !mapped.insert(tensor).second) {
+            valid = false;
+        }
+        mapping.emplace(source, tensor);
+        sources.push_back(source);
+        return tensor;
+    };
+    std::vector<ggml_tensor*> nodes(src->n_nodes), leafs(src->n_leafs);
+    for (int i = 0; i < src->n_nodes; ++i) {
+        if (!(nodes[i] = remap(src->nodes[i])) || !valid) {
+            return false;
+        }
+    }
+    for (int i = 0; i < src->n_leafs; ++i) {
+        if (!(leafs[i] = remap(src->leafs[i])) || !valid) {
+            return false;
+        }
+    }
+    std::vector<std::pair<ggml_tensor*, int32_t>> counts;
+    for (size_t i = 0; i < sources.size(); ++i) {
+        const auto* source = sources[i];
+        if (source->view_src) {
+            remap(source->view_src);
+        }
+        for (const auto* input : source->src) {
+            if (input) {
+                remap(input);
+            }
+        }
+        if (!valid) {
+            return false;
+        }
+        auto* tensor = mapping.at(source);
+        if (!tensor || !src->visited_hash_set.size) {
+            continue;
+        }
+        const size_t index = ggml_hash_find(&src->visited_hash_set, source);
+        if (index == GGML_HASHSET_FULL || !ggml_bitset_get(src->visited_hash_set.used, index)) {
+            continue;
+        }
+        if (counts.size() >= dst->visited_hash_set.size) {
+            return false;
+        }
+        counts.emplace_back(tensor, src->use_counts[index]);
+    }
+    ggml_graph_clear(dst);
+    dst->order = src->order;
+    dst->uid   = 0;
+    for (auto* tensor : nodes) {
+        ggml_graph_add_node(dst, tensor);
+    }
+    for (auto* tensor : leafs) {
+        ggml_graph_add_leaf(dst, tensor);
+    }
+    for (const auto& entry : counts) {
+        const size_t index     = ggml_hash_insert(&dst->visited_hash_set, entry.first);
+        dst->use_counts[index] = entry.second;
+    }
+    return true;
+}
+
+bool ggml_graph_cpy_use_counts(const ggml_cgraph* src, ggml_cgraph* dst) {
+    if (!src || !dst || !src->visited_hash_set.size || !dst->visited_hash_set.size) {
+        return false;
+    }
+    std::unordered_set<ggml_tensor*> represented;
+    for (int i = 0; i < dst->n_nodes; ++i) {
+        represented.insert(dst->nodes[i]);
+    }
+    for (int i = 0; i < dst->n_leafs; ++i) {
+        represented.insert(dst->leafs[i]);
+    }
+    size_t occupied = 0;
+    for (size_t i = 0; i < dst->visited_hash_set.size; ++i) {
+        occupied += ggml_bitset_get(dst->visited_hash_set.used, i);
+    }
+    std::vector<std::pair<ggml_tensor*, int32_t>> counts;
+    for (auto* tensor : represented) {
+        if (!tensor) {
+            return false;
+        }
+        const size_t source_index = ggml_hash_find(&src->visited_hash_set, tensor);
+        if (source_index == GGML_HASHSET_FULL || !ggml_bitset_get(src->visited_hash_set.used, source_index)) {
+            continue;
+        }
+        const size_t index = ggml_hash_find(&dst->visited_hash_set, tensor);
+        if (index == GGML_HASHSET_FULL || !ggml_bitset_get(dst->visited_hash_set.used, index)) {
+            if (++occupied > dst->visited_hash_set.size) {
+                return false;
+            }
+        }
+        counts.emplace_back(tensor, src->use_counts[source_index]);
+    }
+    for (const auto& entry : counts) {
+        const size_t index     = ggml_hash_find_or_insert(&dst->visited_hash_set, entry.first);
+        dst->use_counts[index] = entry.second;
+    }
+    return true;
+}
+
+ggml_status ggml_backend_graph_compute_async_range(ggml_backend_t backend, ggml_cgraph* graph, int first, int last) {
+    if (!backend || !graph || first < 0 || last < first || last > graph->n_nodes) {
+        return GGML_STATUS_FAILED;
+    }
+    if (first == last) {
+        return GGML_STATUS_SUCCESS;
+    }
+    auto view = ggml_graph_view(graph, first, last);
+    return ggml_backend_graph_compute_async(backend, &view);
+}
+#endif
+
+#ifdef SD_USE_UPSTREAM_GGML
+#include "ggml-backend-impl.h"
+#endif
+
+ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void* ptr, size_t size, size_t max_tensor_size) {
+    auto buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
+#ifdef SD_USE_UPSTREAM_GGML
+    if (buffer && !buffer->context) {
+        auto reg = ggml_backend_dev_backend_reg(device);
+        if (reg && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
+            // Older Metal backends wrap failed mappings and cannot free their null context.
+            buffer->iface.free_buffer = nullptr;
+            ggml_backend_buffer_free(buffer);
+            return nullptr;
+        }
+    }
+#endif
+    return buffer;
+}