change use_counts to be indexed by hash table slot

Author: jeffbolznv

ggml/src/ggml-backend.cpp

@@ -818,7 +818,8 @@ static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, str
         if (sched->debug > 1) {
             ggml_backend_t tensor_backend = ggml_backend_sched_get_tensor_backend(sched, node);
             GGML_LOG_DEBUG("node #%3d (%10.10s): %20.20s (%5.5s) [%5.5s %8.8s] use=%d:", i, ggml_op_name(node->op), node->name,
-                fmt_size(ggml_nbytes(node)), tensor_backend ? ggml_backend_name(tensor_backend) : "NULL", GET_CAUSE(node), graph->use_counts[i]);
+                fmt_size(ggml_nbytes(node)), tensor_backend ? ggml_backend_name(tensor_backend) : "NULL", GET_CAUSE(node),
+                graph->use_counts[ggml_hash_find(&graph->visited_hash_set, node)]);
             for (int j = 0; j < GGML_MAX_SRC; j++) {
                 struct ggml_tensor * src = node->src[j];
                 if (src == NULL) {

ggml/src/ggml-impl.h

@@ -301,8 +301,7 @@ struct ggml_cgraph {
     struct ggml_tensor ** grads;     // the outputs of these tensors are the gradients of the nodes
     struct ggml_tensor ** grad_accs; // accumulators for node gradients
     struct ggml_tensor ** leafs;     // tensors with constant data
-    int32_t             * use_counts;// number of uses of each tensor
-    int32_t             * hash_to_node;// map hash index to node index
+    int32_t             * use_counts;// number of uses of each tensor, indexed by hash table slot
 
     struct ggml_hash_set visited_hash_set;
 
@@ -474,7 +473,7 @@ static inline ggml_bf16_t ggml_compute_fp32_to_bf16(float s) {
 static inline bool ggml_node_has_N_uses(const struct ggml_cgraph * cgraph, int node_idx, int32_t N) {
     const struct ggml_tensor * node = cgraph->nodes[node_idx];
     // check the use count against how many we're replacing
-    if (cgraph->use_counts[node_idx] != N) {
+    if (cgraph->use_counts[ggml_hash_find(&cgraph->visited_hash_set, node)] != N) {
         return false;
     }
 

ggml/src/ggml.c

@@ -5809,8 +5809,7 @@ static void ggml_visit_parents(struct ggml_cgraph * cgraph, struct ggml_tensor *
     if (ggml_hash_insert(&cgraph->visited_hash_set, node) == GGML_HASHSET_ALREADY_EXISTS) {
         return;
     }
-
-    bool incr_use_count = false;
+    cgraph->use_counts[ggml_hash_find(&cgraph->visited_hash_set, node)] = 0;
 
     for (int i = 0; i < GGML_MAX_SRC; ++i) {
         const int k =
@@ -5825,14 +5824,7 @@ static void ggml_visit_parents(struct ggml_cgraph * cgraph, struct ggml_tensor *
             // Update the use count for this operand.
             // Skip if it's a leaf node
             if (!(s->op == GGML_OP_NONE && !(s->flags & GGML_TENSOR_FLAG_PARAM))) {
-                // the src can be the node itself (happens in ggml_cast)
-                if (s == node) {
-                    incr_use_count = true;
-                } else {
-                    size_t s_idx = cgraph->hash_to_node[ggml_hash_find(&cgraph->visited_hash_set, s)];
-                    GGML_ASSERT(cgraph->nodes[s_idx] == s);
-                    cgraph->use_counts[s_idx]++;
-                }
+                cgraph->use_counts[ggml_hash_find(&cgraph->visited_hash_set, s)]++;
             }
         }
     }
@@ -5855,10 +5847,6 @@ static void ggml_visit_parents(struct ggml_cgraph * cgraph, struct ggml_tensor *
         }
 
         cgraph->nodes[cgraph->n_nodes] = node;
-        if (incr_use_count) {
-            cgraph->use_counts[cgraph->n_nodes]++;
-        }
-        cgraph->hash_to_node[ggml_hash_find(&cgraph->visited_hash_set, node)] = cgraph->n_nodes;
         cgraph->n_nodes++;
     }
 }
@@ -5995,11 +5983,10 @@ static void * incr_ptr_aligned(void ** p, size_t size, size_t align) {
 static size_t ggml_graph_nbytes(size_t size, bool grads) {
     size_t hash_size = ggml_hash_size(size * 2);
     void * p = 0;
-    incr_ptr_aligned(&p,             sizeof(struct ggml_cgraph),   1);
-    incr_ptr_aligned(&p, size      * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // nodes
-    incr_ptr_aligned(&p, size      * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // leafs
-    incr_ptr_aligned(&p, size      * sizeof(int32_t),              sizeof(int32_t)); // use_counts
-    incr_ptr_aligned(&p, hash_size * sizeof(int32_t),              sizeof(int32_t)); // hash_to_node
+    incr_ptr_aligned(&p, sizeof(struct ggml_cgraph), 1);
+    incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // nodes
+    incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // leafs
+    incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); // use_counts
     incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // hash keys
     if (grads) {
         incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grads
@@ -6031,8 +6018,7 @@ struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t siz
 
     struct ggml_tensor ** nodes_ptr     =         incr_ptr_aligned(&p, size      * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *));
     struct ggml_tensor ** leafs_ptr     =         incr_ptr_aligned(&p, size      * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *));
-    int32_t             * use_counts_ptr=         incr_ptr_aligned(&p, size      * sizeof(int32_t), sizeof(int32_t));
-    int32_t             * hash_to_node_ptr=       incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t));
+    int32_t             * use_counts_ptr=         incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t));
     struct ggml_tensor ** hash_keys_ptr =         incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *));
     struct ggml_tensor ** grads_ptr     = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL;
     struct ggml_tensor ** grad_accs_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL;
@@ -6051,7 +6037,6 @@ struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t siz
         /*.grad_accs    =*/ grad_accs_ptr,
         /*.leafs        =*/ leafs_ptr,
         /*.use_counts   =*/ use_counts_ptr,
-        /*.hash_to_node =*/ hash_to_node_ptr,
         /*.hash_table   =*/ { hash_size, hash_used, hash_keys_ptr },
         /*.order        =*/ GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT,
     };
@@ -6061,8 +6046,6 @@ struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t siz
         memset(cgraph->grads,     0, hash_size*sizeof(struct ggml_tensor *));
         memset(cgraph->grad_accs, 0, hash_size*sizeof(struct ggml_tensor *));
     }
-    memset(cgraph->use_counts,    0, size*sizeof(int32_t));
-    memset(cgraph->hash_to_node, -1, size*sizeof(int32_t));
 
     return cgraph;
 }
@@ -6080,9 +6063,8 @@ struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph0, int i0, int i1)
         /*.grads            =*/ NULL, // gradients would need visited_hash_set
         /*.grad_accs        =*/ NULL,
         /*.leafs            =*/ NULL,
-        /*.use_counts       =*/ cgraph0->use_counts + i0,
-        /*.hash_to_node     =*/ NULL,
-        /*.visited_hash_set =*/ { 0, NULL, NULL },
+        /*.use_counts       =*/ cgraph0->use_counts,
+        /*.visited_hash_set =*/ cgraph0->visited_hash_set,
         /*.order            =*/ cgraph0->order,
     };
 
@@ -6104,14 +6086,13 @@ void ggml_graph_cpy(struct ggml_cgraph * src, struct ggml_cgraph * dst) {
 
     for (int i = 0; i < src->n_nodes; ++i) {
         dst->nodes[i] = src->nodes[i];
-        dst->use_counts[i] = src->use_counts[i];
     }
 
     for (size_t i = 0; i < src->visited_hash_set.size; ++i) {
         // copy all hashset keys (tensors) that are in use
         if (ggml_bitset_get(src->visited_hash_set.used, i)) {
             ggml_hash_insert(&dst->visited_hash_set, src->visited_hash_set.keys[i]);
-            dst->hash_to_node[i] = src->hash_to_node[i];
+            dst->use_counts[i] = src->use_counts[i];
         }
     }