Memory allocation method, related device, and computer-readable storage medium

Abstract

This application provides a memory allocation method. The method includes: obtaining a computation graph corresponding to a neural network; sequentially allocating memory space to M pieces of tensor data based on a sorting result of the M pieces of tensor data, where if at least a part of the allocated memory space can be reused for one of the M pieces of tensor data, the at least a part of the memory space that can be reused for the tensor data is allocated to the tensor data, the allocated memory space is memory space that has been allocated to the M pieces of tensor data before the tensor data, the sorting result indicates a sequence of allocating memory space to the M pieces of tensor data, and the sorting result is related to information about each of the M pieces of tensor data.

1 . A memory allocation method, comprising:
 obtaining a computation graph corresponding to a neural network, wherein the computation graph comprises N nodes and a directed edge that connects different nodes, the directed edge of the computation graph carrying a piece of tensor data of M pieces of tensor data, the computation graph comprising the M pieces of tensor data, wherein M is an integer greater than 1; and   sequentially allocating memory space to the M pieces of tensor data based on a sorting result of the M pieces of tensor data, wherein in response to that at least one part of allocated memory space can be reused for one piece of tensor data of the M pieces of tensor data, the at least one part of the memory space that can be reused for the one piece of tensor data is allocated to the piece of tensor data, wherein the allocated memory space is memory space that has been allocated to at least one piece of tensor data of the M pieces of tensor data before allocating memory space to the piece of tensor data, wherein the sorting result indicates a sequence of allocating memory space to the M pieces of tensor data, wherein the sorting result is related to information about each piece of tensor data of the M pieces of tensor data, wherein the information about each piece of tensor data indicates at least one of: a constraint relationship corresponding to each piece of tensor data and a quantity of nodes to which each piece of tensor data flows, or the constraint relationship indicating a relationship between available memory space of one piece of tensor data of the M pieces of tensor data and available memory space of other one or more pieces of tensor data in the M pieces of tensor data.   
     
     
         2 . The method according to  claim 1 , wherein in response to that the allocated memory space cannot be reused for the piece of tensor data, other memory space is allocated to the piece of tensor data, wherein the other memory space is different from the allocated memory space. 
     
     
         3 . The method according to  claim 1 , wherein the constraint relationship indicates at least one of: a relationship between available memory space of one piece of tensor data and available memory space of another piece of tensor data is reusable, the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is non-reusable, or the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is non-reusable and continuous. 
     
     
         4 . The method according to  claim 3 , wherein the constraint relationship is carried in a constraint relationship table; the constraint relationship table comprising identifiers of the M pieces of tensor data; and wherein in the constraint relationship table, a first value indicates that the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is reusable, a second value indicates that the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is non-reusable, and a third value indicates that the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is non-reusable and continuous. 
     
     
         5 . The method according to  claim 1 , wherein in response to that all consumption nodes of first tensor data are upstream nodes of a production node of second tensor data, or in response to that all consumption nodes of the second tensor data are downstream nodes of a production node of the first tensor data, memory space allocated to the second tensor data can be reused for the first tensor data; and wherein in response to that all the consumption nodes of the first tensor data are not the upstream nodes of the production node of the second tensor data, or in response to that all the consumption nodes of the second tensor data are not the downstream nodes of the production node of the first tensor data, the memory space allocated to the second tensor data cannot be reused for the first tensor data, wherein the first tensor data and the second tensor data are any two pieces of tensor data of the M pieces of tensor data, wherein the consumption node is a node to which tensor data flows, and the production node is a node from which tensor data flows. 
     
     
         6 . The method according to  claim 1 , wherein the computation graph comprises a plurality of computing subtasks, wherein the computing subtask indicates a computing function by using a group of nodes and an edge related to the group of nodes, and wherein an execution relationship between the plurality of computing subtasks is parallel; and wherein the method further comprises:
 in one computing subtask, in response to that there is no directed edge between two adjacent nodes, adding a directed edge between the two adjacent nodes, to update the computation graph to obtain an updated computation graph, wherein each added directed edge carries corresponding piece of tensor data, and wherein the two adjacent nodes are two nodes that are adjacent in an execution sequence in the computing subtask; and   obtaining the information about each piece of tensor data of the M pieces of tensor data based on the updated computation graph.   
     
     
         7 . The method according to  claim 6 , wherein the computation graph further comprises a first computing subtask and a second computing subtask that are in a serial execution relationship, and wherein the first computing subtask is before the second computing subtask in an execution sequence; and wherein the updating the computation graph further comprises:
 in response to that there is no directed edge between a last node in the first computing subtask and a first node in the second computing subtask, adding a directed edge between the last node in the first computing subtask and the first node in the second computing subtask.   
     
     
         8 . The method according to  claim 1 , wherein in the computation graph, a value of an identifier of a production node of a piece of tensor data is less than a value of an identifier of a consumption node of the piece of tensor data, and the production node of the tensor data and the consumption node of the piece of tensor data are two adjacent nodes. 
     
     
         9 . The method according to  claim 8 , wherein an identifier of each node in the computation graph is used to determine the information about each piece of tensor data of the M pieces of tensor data. 
     
     
         10 . The method according to  claim 1 , wherein the information about each piece of tensor data indicates the constraint relationship corresponding to each piece of tensor data, the method further comprising:
 obtaining, based on the constraint relationship corresponding to each piece of tensor data, constraint amounts respectively corresponding to the M pieces of tensor data, wherein each constraint amount of the constraint amounts is an amount of pieces of tensor data of the M pieces of tensor data for which same memory space cannot be reused with a piece of tensor data of the M pieces of tensor data; and   sorting the M pieces of tensor data based on the constraint amounts respectively corresponding to the M pieces of tensor data, to obtain the sorting result of the M pieces of tensor data.   
     
     
         11 . The method according to  claim 1 , wherein the information about each piece of tensor data indicates a quantity of nodes to which each piece of tensor data flows, the method further comprising:
 sorting the M pieces of tensor data based on quantities of consumption nodes that respectively correspond to the M pieces of tensor data, to obtain the sorting result of the M pieces of tensor data.   
     
     
         12 . The method according to  claim 1 , the method further comprising:
 sorting the M pieces of tensor data based on the information about each piece of tensor data according to a heuristic algorithm, to obtain the sorting result of the M pieces of tensor data within a preset time period.   
     
     
         13 . The method according to  claim 12 , wherein the sorting result is a sorting result obtained after an optimization, and wherein a size of maximum memory that needs to be occupied by the neural network and that corresponds to the sorting result obtained after the optimization is less than a size of maximum memory that needs to be occupied by the neural network and that is determined based on a sorting result existing before the optimization. 
     
     
         14 . A memory allocation apparatus, comprising:
 at least one processor; and   at least one processor memory coupled to the at least one processor to store program instructions, which when executed by the processor, cause the at least one processor to:   obtain a computation graph corresponding to a neural network, wherein the computation graph comprises N nodes and a directed edge that connects different nodes, the directed edge of the computation graph carrying a piece of tensor data of M pieces of tensor data, the computation graph comprising the M pieces of tensor data, wherein M is an integer greater than 1; and   sequential allocate memory space to the M pieces of tensor data based on a sorting result of the M pieces of tensor data, wherein in response to that at least one part of allocated memory space can be reused for one piece of tensor data of the M pieces of tensor data, the at least one part of the memory space that can be reused for the one piece of tensor data is allocated to the piece of tensor data, wherein the allocated memory space is memory space that has been allocated to at least one piece of tensor data of the M pieces of tensor data before allocating memory space to the piece of tensor data, the sorting result indicates a sequence of allocating memory space to the M pieces of tensor data, wherein the sorting result is related to information about each piece of tensor data of the M pieces of tensor data, wherein the information about each piece of tensor data indicates at least one of: a constraint relationship corresponding to each piece of tensor data and a quantity of nodes to which each piece of tensor data flows, or the constraint relationship indicating a relationship between available memory space of one piece of tensor data of the M pieces of tensor data and available memory space of other one or more pieces of tensor data in the M pieces of tensor data.   
     
     
         15 . The apparatus according to  claim 14 , wherein in response to that the allocated memory space cannot be reused for the piece of tensor data, other memory space is allocated to the piece of tensor data, wherein the other memory space is different from the allocated memory space. 
     
     
         16 . The apparatus according to  claim 14 , wherein the constraint relationship indicates at least one of: a relationship between available memory space of one piece of tensor data and available memory space of another piece of tensor data is reusable, the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is non-reusable, or the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is non-reusable and continuous. 
     
     
         17 . The apparatus according to  claim 16 , wherein the constraint relationship is carried in a constraint relationship table; the constraint relationship table comprising identifiers of the M pieces of tensor data; and wherein in the constraint relationship table, a first value indicates that the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is reusable, a second value indicates that the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is non-reusable, and a third value indicates that the relationship between the available memory space of the one piece of tensor data and the available memory space of the another piece of tensor data is non-reusable and continuous. 
     
     
         18 . The apparatus according to  claim 14 , wherein in response to that all consumption nodes of first tensor data are upstream nodes of a production node of second tensor data, or in response to that all consumption nodes of the second tensor data are downstream nodes of a production node of the first tensor data, memory space allocated to the second tensor data can be reused for the first tensor data; and wherein in response to that all the consumption nodes of the first tensor data are not the upstream nodes of the production node of the second tensor data, or in response to that all the consumption nodes of the second tensor data are not the downstream nodes of the production node of the first tensor data, the memory space allocated to the second tensor data cannot be reused for the first tensor data, wherein the first tensor data and the second tensor data are any two pieces of tensor data of the M pieces of tensor data, the consumption node is a node to which tensor data flows, and wherein the production node is a node from which tensor data flows. 
     
     
         19 . The apparatus according to  claim 14 , wherein the computation graph comprises a plurality of computing subtasks, wherein the computing subtask indicates a computing function by using a group of nodes and an edge related to the group of nodes, and wherein an execution relationship between the plurality of computing subtasks is parallel; and wherein the at least one processor is further to:
 in one computing subtask, in response to that there is no directed edge between two adjacent nodes, add a directed edge between the two adjacent nodes, to update the computation graph to obtain an updated computation graph, wherein each added directed edge carries corresponding piece of tensor data, and wherein the two adjacent nodes are two nodes that are adjacent in an execution sequence in the computing subtask; and   obtain the information about each piece of tensor data of the M pieces of tensor data based on the updated computation graph.   
     
     
         20 . AA-non-transitory computer-readable storage medium, storing one or more instructions that, when executed by at least one processor, cause the at least one processor to:
 obtain a computation graph corresponding to a neural network, wherein the computation graph comprises N nodes and a directed edge that connects different nodes, the directed edge of the computation graph carrying a piece of tensor data of M pieces of tensor data, the computation graph comprising the M pieces of tensor data, and wherein M is an integer greater than 1; and   sequentially allocate memory space to the M pieces of tensor data based on a sorting result of the M pieces of tensor data, wherein in response to that at least one part of the allocated memory space can be reused for one piece of tensor data of the M pieces of tensor data, the at least one part of the memory space that can be reused for the one piece of tensor data is allocated to the piece of tensor data, wherein the allocated memory space is memory space that has been allocated to at least one piece of tensor data of the M pieces of tensor data before the tensor data, wherein the sorting result indicates a sequence of allocating memory space to the M pieces of tensor data, wherein the sorting result is related to information about each piece of tensor data of the M pieces of tensor data, wherein the information about each piece of tensor data indicates at least one of: a constraint relationship corresponding to each piece of tensor data and a quantity of nodes to which each piece of tensor data flows, or the constraint relationship indicating a relationship between available memory space of one piece of tensor data of the M pieces of tensor data and available memory space of other one or more pieces of tensor data in the M pieces of tensor data.