Skip to content

Commit 58b6cb6

Browse files
authored
fix: prevent gallocr hash overflow in tiny graph-cut segments (#1880)
1 parent 6100d83 commit 58b6cb6

1 file changed

Lines changed: 39 additions & 7 deletions

File tree

src/core/ggml_graph_cut.cpp

Lines changed: 39 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -603,7 +603,43 @@ namespace sd::ggml_graph_cut {
603603
GGML_ASSERT(gf != nullptr);
604604
GGML_ASSERT(graph_ctx_out != nullptr);
605605

606-
const size_t graph_size = segment.internal_node_indices.size() + segment.input_refs.size() + 8;
606+
// Collect leaf inputs and internal nodes, then any tensor they
607+
// reference that is not already represented, notably the view_src of a
608+
// view-typed input leaf. ggml_gallocr sizes its hash set from
609+
// n_nodes + n_leafs (plus a 25% margin that rounds down to zero for a
610+
// one-node segment), so every distinct tensor it will hash must be
611+
// counted here or a tiny segment overflows the hash set and aborts.
612+
std::vector<ggml_tensor*> leaves;
613+
std::unordered_set<ggml_tensor*> represented;
614+
for (const auto& input : segment.input_refs) {
615+
ggml_tensor* current_input = input_tensor(gf, input);
616+
if (current_input == nullptr) {
617+
continue;
618+
}
619+
if (represented.insert(current_input).second) {
620+
leaves.push_back(current_input);
621+
}
622+
}
623+
for (int node_idx : segment.internal_node_indices) {
624+
represented.insert(ggml_graph_node(gf, node_idx));
625+
}
626+
auto add_reference = [&](ggml_tensor* tensor) {
627+
if (tensor != nullptr && represented.insert(tensor).second) {
628+
leaves.push_back(tensor);
629+
}
630+
};
631+
for (int node_idx : segment.internal_node_indices) {
632+
ggml_tensor* node = ggml_graph_node(gf, node_idx);
633+
for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
634+
add_reference(node->src[src_idx]);
635+
}
636+
add_reference(node->view_src);
637+
}
638+
for (size_t i = 0; i < leaves.size(); ++i) {
639+
add_reference(leaves[i]->view_src);
640+
}
641+
642+
const size_t graph_size = segment.internal_node_indices.size() + leaves.size() + 8;
607643
ggml_init_params params = {
608644
/*.mem_size =*/ggml_graph_overhead_custom(graph_size, false) + 1024,
609645
/*.mem_buffer =*/nullptr,
@@ -614,13 +650,9 @@ namespace sd::ggml_graph_cut {
614650
ggml_cgraph* segment_graph = ggml_new_graph_custom(graph_ctx, graph_size, false);
615651
GGML_ASSERT(segment_graph != nullptr);
616652

617-
for (const auto& input : segment.input_refs) {
618-
ggml_tensor* current_input = input_tensor(gf, input);
619-
if (current_input == nullptr) {
620-
continue;
621-
}
653+
for (ggml_tensor* leaf : leaves) {
622654
GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
623-
segment_graph->leafs[segment_graph->n_leafs++] = current_input;
655+
segment_graph->leafs[segment_graph->n_leafs++] = leaf;
624656
}
625657

626658
for (int output_node_index : segment.output_node_indices) {

0 commit comments

Comments
 (0)