@@ -610,6 +610,26 @@ static bool ggml_gallocr_is_allocated(ggml_gallocr_t galloc, struct ggml_tensor
610610 return t -> data != NULL || ggml_gallocr_hash_get (galloc , t )-> allocated ;
611611}
612612
613+ // free the extra space at the end if the new tensor is smaller
614+ static void ggml_gallocr_free_extra_space (ggml_gallocr_t galloc , struct ggml_tensor * node , struct ggml_tensor * parent ) {
615+ struct hash_node * hn = ggml_gallocr_hash_get (galloc , node );
616+ struct hash_node * p_hn = ggml_gallocr_hash_get (galloc , parent );
617+
618+ size_t parent_size = ggml_backend_buft_get_alloc_size (galloc -> bufts [p_hn -> buffer_id ], parent );
619+ size_t node_size = ggml_backend_buft_get_alloc_size (galloc -> bufts [hn -> buffer_id ], node );
620+
621+ GGML_ASSERT (parent_size >= node_size );
622+
623+ if (parent_size > node_size ) {
624+ struct ggml_dyn_tallocr * p_alloc = galloc -> buf_tallocs [p_hn -> buffer_id ];
625+ struct buffer_address p_addr = p_hn -> addr ;
626+ p_addr .offset += node_size ;
627+ size_t extra_size = parent_size - node_size ;
628+ AT_PRINTF ("freeing extra %zu bytes from parent %s for %s\n" , extra_size , parent -> name , node -> name );
629+ ggml_dyn_tallocr_free_tensor (p_alloc , p_addr , extra_size , parent );
630+ }
631+ }
632+
613633static void ggml_gallocr_allocate_node (ggml_gallocr_t galloc , struct ggml_tensor * node , int buffer_id ) {
614634 struct hash_node * hn = ggml_gallocr_hash_get (galloc , node );
615635
@@ -654,13 +674,15 @@ static void ggml_gallocr_allocate_node(ggml_gallocr_t galloc, struct ggml_tensor
654674 hn -> addr = p_hn -> addr ;
655675 p_hn -> allocated = false; // avoid freeing the parent
656676 view_src_hn -> allocated = false;
677+ ggml_gallocr_free_extra_space (galloc , node , view_src );
657678 return ;
658679 }
659680 } else {
660681 AT_PRINTF ("reusing parent %s for %s\n" , parent -> name , node -> name );
661682 hn -> buffer_id = p_hn -> buffer_id ;
662683 hn -> addr = p_hn -> addr ;
663684 p_hn -> allocated = false; // avoid freeing the parent
685+ ggml_gallocr_free_extra_space (galloc , node , parent );
664686 return ;
665687 }
666688 }
0 commit comments