Executing computational graphs on graphics processing units
Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for generating a data entity that causes a processing unit to process a computational graph. In one aspect, a method includes the actions of receiving data identifying a computational graph, the computational graph including a plurality of nodes representing operations; obtaining compilation artifacts for processing the computational graph on a processing unit; and generating a data entity from the compilation artifacts, wherein the data entity, when invoked, causes the processing unit to process the computational graph by executing the operations represented by the plurality of nodes.
1. A method for compiling machine code for one or more processing units that are controlled by a processor, the method comprising:
receiving a computational graph and input buffer parameters, wherein the computational graph comprises a plurality of nodes representing respective operations;
generating compilation artifacts, wherein the compilation artifacts include the machine code for execution by the one or more processing units to execute the respective operations represented by the plurality of nodes in the computational graph, wherein the compilation artifacts include the input buffer parameters, and wherein the compilation artifacts exclude any machine code for execution by the processor;
generating, from the compilation artifacts, a data entity that, when invoked at run time by the one or more processing units, causes the one or more processing units to process the computational graph in accordance with the input buffer parameters and without executing any machine code by the processor, the processing by the one or more processing units comprising:
enqueuing the respective operations into the one or more processing units,
based on the data entity, defining dependencies between the respective operations, and
executing the respective enqueued operations by the one or more processing units according to the dependencies and without executing any machine code by the processor.
2. The method of claim 1 , wherein the one or more processing units comprise a graphics processing unit (GPU) and the processor comprises a central processing unit (CPU).
3. The method of claim 1 , wherein the data entity includes machine code that, when executed by the one or more processing units, cause the one or more processing units to invoke the data entity to process the computational graph in accordance with the input buffer parameters and without awaiting the instruction data from the processor.
4. The method of claim 1 , wherein the respective operations are enqueued according to a particular order; and the processing comprises:
generating instructions that when executed by the one or more processing units cause the one or more processing units to execute the respective operations according to the particular order.
5. The method of claim 4 , wherein enqueuing the respective operations into the one or more processing units, comprises:
generating instructions that when executed by the one or more processing units cause the one or more processing units to enqueue, at each processing unit in the one or more processing units, a respective subset of the respective operations; and
generating instructions that when executed by the one or more processing units cause the one or more processing units to, at each processing unit, execute the respective subset of the respective operations for the processing unit.
6. The method of claim 1 , wherein each processing unit of the one or more processing units comprises a respective buffer, and the method further comprises:
assigning, by the one or more processing units, a first operation of the respective operations to a buffer of a particular processing unit of the one or more processing units for execution based on buffer data.
7. The method of claim 1 , wherein the input buffer parameters comprise:
a plurality of buffer parameters including the input buffer parameters and associations between the plurality of buffer parameters and the respective operations, each of the plurality of buffer parameters being associated with a respective operation of the respective operations.
8. The method of claim 1 , wherein the compilation artifacts comprise:
a data structure representing (i) a sequence of operations corresponding to the respective operations represented by the computational graph, and (ii) dependencies between the respective operations.
9. The method of claim 1 , wherein the compilation artifacts comprise:
library data representing a plurality of buffer parameters and associations between the plurality of buffer parameters and a plurality of libraries comprising one or more sub-routines, each of the plurality of buffer parameters associated with a respective library of the plurality of libraries.
10. The method of claim 1 , wherein the respective operations are operations for processing an input of a neural network through one or more layers of the neural network to generate an output of the neural network.
11. A system comprising:
one or more computers and one or more storage devices on which are stored instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform first operations for compiling machine code for one or more processing units, comprising:
receiving a computational graph and input buffer parameters, wherein the computational graph comprises a plurality of nodes representing respective operations;
generating compilation artifacts, wherein the compilation artifacts include the machine code for execution by the one or more processing units to execute the respective operations represented by the plurality of nodes in the computational graph, wherein the compilation artifacts include the input buffer parameters, and wherein the compilation artifacts exclude any machine code for execution by a processor, the processor configured to control the one or more processing units;
generating, from the compilation artifacts, a data entity that when invoked at run time by the one or more processing units, causes the one or more processing units to process the computational graph in accordance with the input buffer parameters and without executing any machine code by the processor, the processing by the one or more processing units comprising:
enqueuing the respective operations into the one or more processing units,
based on the data entity, defining dependencies between the respective operations, and
executing the respective enqueued operations by the one or more processing units according to the dependencies and without executing any machine code by the processor.
12. The system of claim 11 , wherein the one or more processing units comprise a graphics processing unit (GPU) and the processor comprises a central processing unit (CPU).
13. The system of claim 11 , wherein the data entity includes machine code that, when executed by the one or more processing units, causes the one or more processing units to invoke the data entity to process the computational graph in accordance with the input buffer parameters and without awaiting the instruction data from the processor.
14. The system of claim 11 , wherein the respective operations are enqueued according to a particular order; and the processing further comprises:
generating instructions that when executed by the one or more processing units causes the one or more processing units to execute the respective operations according to the particular order.
15. The system of claim 14 , wherein enqueuing the respective operations into the one or more processing units, comprises:
generating instructions that when executed by the one or more processing units causes the one or more processing units to enqueue, at each processing unit in the one or more processing units, a respective subset of the respective operations; and
generating instructions that when executed by the one or more processing units causes the one or more processing units to, at each processing unit, execute the respective subset of the respective operations for the processing unit.
16. The system of claim 11 , wherein each processing unit of the one or more processing units comprises a respective buffer, and the first operations further comprising:
assigning, by the one or more processing units, a first second operation of the respective operations to a buffer of a particular processing unit of the one or more processing units for execution based on buffer data.
17. The system of claim 11 , wherein the input buffer parameters comprise:
a plurality of buffer parameters including the input buffer parameters and associations between the plurality of buffer parameters and the respective operations, each of the plurality of buffer parameters being associated with a respective second operation of the respective operations.
18. The system of claim 11 , wherein the compilation artifacts comprise:
library data representing a plurality of buffer parameters and associations between the plurality of buffer parameters and a plurality of libraries comprising one or more sub-routines, each of the plurality of buffer parameters associated with a respective library of the plurality of libraries.
19. The system of claim 11 , wherein the respective operations are operations for processing an input of a neural network through one or more layers of the neural network to generate an output of the neural network.
20. One or more non-transitory computer-readable storage media encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform first operations for compiling machine code for one or more processing units, the first operations comprising:
receiving a computational graph and input buffer parameters, wherein the computational graph comprises a plurality of nodes representing respective operations;
generating compilation artifacts, wherein the compilation artifacts include the machine code for execution by the one or more processing units to execute the respective operations represented by the plurality of nodes in the computational graph, wherein the compilation artifacts include the input buffer parameters, and wherein the compilation artifacts exclude any machine code for execution by a processor, the processor configured to control the one or more processing units;
generating, from the compilation artifacts, a data entity that when invoked at run time by the one or more processing units, causes the one or more processing units to process the computational graph in accordance with the input buffer parameters by executing the respective operations and without executing any machine code by the processor, the processing by the one or more processing units comprising:
enqueuing the respective operations into the one or more processing units,
based on the data entity, defining dependencies between the respective operations, and
executing the respective enqueued operations by the one or more processing units according to the dependencies and without executing any machine code by the processor.