Stream-based accelerator processing of computational graphs

1. A method for distributing a computation represented by a computational graph over a plurality of devices including a stream-based accelerator device and performing the computation including performing the computation represented by a subgraph of the computational graph on streams of the stream-based accelerator device, the method comprising:

receiving, by a computational graph system, a request to perform a computation represented by a computational graph;

obtaining data representing multiple subgraphs of the computational graph, the computational graph comprising a plurality of nodes and directed edges, wherein each node represents a respective operation, wherein each directed edge connects a respective first node to a respective second node that represents an operation that receives, as input, an output of an operation represented by the respective first node;

assigning the multiple subgraphs to respective devices of a plurality of devices, including assigning a first subgraph to a first device by a placer in the computational graph system;

determining that the first device comprises a hardware accelerator having a plurality of streams, wherein each stream is an independent hardware queue that performs operations assigned to the stream in the order assigned;

in response to determining that the first device comprises a hardware accelerator having a plurality of streams, generating instructions that when executed by the first device cause the first device to:

assign the operation represented by each node in the first subgraph to a respective stream in the plurality of streams of the hardware accelerator, including identifying two operations that do not depend on each other and assigning the two operations to different streams, and assigning to a single stream a first operation that changes an internal state of the hardware in a way that must happen before a second operation executes and assigning the second operation to be performed in the single stream after the first operation is complete; and

perform the operations represented by the nodes in the first subgraph in accordance with the assignment;

providing the instructions and the data to the first device;

sending to the first device a request causing the first device to perform the operations represented by the nodes in the first subgraph;

obtaining from the first device outputs of operations performed by the first device;

causing the other devices of the plurality of devices to perform the operations of the subgraphs assigned to the other devices; and

responding to the request to perform the computation with one or more outputs obtained from the plurality of devices.

2. The method of claim 1 , wherein the request specifies identifying one or more particular outputs from one or more respective nodes in the subgraph, further comprising:

receiving, from the first device, the one or more particular outputs; and

providing the one or more particular outputs to the client.

3. The method of claim 1 , wherein the instructions further cause the first device to store the one or more particular outputs in memory of the first device.

4. The method of claim 1 , wherein the operations for the first subgraph comprise partial inference or training computations for a neural network.

5. The method of claim 1 , further comprising:

analyzing the first subgraph to identify a group of nodes in the first subgraph in a chain structure;

wherein the instructions cause the first device to assign the group of nodes to one stream.

6. The method of claim 1 , wherein the assigning comprises:

analyzing the first subgraph to identify a first node in the first subgraph having a plurality of directed edges as outputs;

wherein the instructions cause the first device to assign, for each of the directed edges, a node to which the directed edge points to a disjoint stream of the hardware accelerator.

7. The method of claim 1 , wherein the instructions cause the first device to determine, for each node, a respective amount of memory resources in the hardware accelerator consumed by the operation represented by the node based on the directed edges to the node, wherein the assigning is based at least on the respective amount of memory resources.

8. The method of claim 1 , wherein the instructions cause the first device to determine a particular operation represented by a node has finished at a particular stream;

in response to determining the particular operation has finished:

determine a first amount of memory consumed by the particular operation that will be freed;

determine, for each of a group of unassigned nodes, a respective estimated amount of memory consumed by an operation that is represented by the unassigned node;

determine, from the group of unassigned nodes, a first unassigned node that represents an operation, which executes on a stream of the hardware accelerator, with the estimated amount of memory that maximizes usage of the first amount of memory; and

assign an operation represented by the first unassigned node to the particular stream.

9. The method of claim 1 , wherein the instructions cause the first device to determine a particular operation represented by a node has finished at a particular stream:

in response to determining the particular operation has finished:

determine at least one subsequent operation that uses the output of the particular operation as input; and

reuse memory allocated for the output of the particular operation after the at least one subsequent operation has executed.

10. The method of claim 9 , wherein determining at least one subsequent operation that uses the output of the particular operation as input includes:

determining that at least two subsequent operations, a first subsequent operation in a first stream and a second subsequent operation in a second stream, use the output of the particular operation as input;

placing a first marker in a first stream that indicates when the first subsequent operation has used the particular operation as input;

placing a second marker in a second stream that indicates when the second subsequent operation has used the particular operation as input;

determining that both operations have used the particular operation upon indication from the first and second markers.

11. A system for distributing a computation represented by a computational graph over a plurality of devices including a stream-based accelerator device and performing the computation including performing the computation represented by a subgraph of the computational graph on streams of the stream-based accelerator device, the system comprising:

one or more computers; and

computer-readable medium coupled to the one or more computers and having instructions stored thereon, which, when executed by the one or more computers, cause the one or more computers to perform operations comprising:

receiving, by a computational graph system, a request to perform a computation represented by a computational graph;

obtaining data representing multiple subgraphs of the computational graph, the computational graph comprising a plurality of nodes and directed edges, wherein each node represents a respective operation, wherein each directed edge connects a respective first node to a respective second node that represents an operation that receives, as input, an output of an operation represented by the respective first node;

assigning the multiple subgraphs to respective devices of a plurality of devices, including assigning a first subgraph to a first device by a placer in the computational graph system;

determining that the first device comprises a hardware accelerator having a plurality of streams, wherein each stream is an independent hardware queue that performs operations assigned to the stream in the order assigned;

in response to determining that the first device comprises a hardware accelerator having a plurality of streams, generating instructions that when executed by the first device cause the first device to:

assign the operation represented by each node in the first subgraph to a respective stream in the plurality of streams of the hardware accelerator, including identifying two operations that do not depend on each other and assigning the two operations to different streams, and assigning to a single stream a first operation that changes an internal state of the hardware in a way that must happen before a second operation executes and assigning the second operation to be performed in the single stream after the first operation is complete; and

perform the operations represented by the nodes in the subgraph in accordance with the assignment;

providing the instructions and the data to the first device;

sending to the first device a request causing the first device to perform the operations represented by the nodes in the first subgraph;

obtaining from the first device outputs of operations performed by the first device;

causing the other devices of the plurality of devices to perform the operations of the subgraphs assigned to the other devices; and

responding to the request to perform the computation with one or more outputs obtained from the plurality of devices.

12. The system of claim 11 , wherein the request specifies identifying one or more particular outputs from one or more respective nodes in the subgraph, further comprising:

receiving, from the first device, the one or more particular outputs; and

providing the one or more particular outputs to the client.

13. The system of claim 11 , further comprising:

analyzing the first subgraph to identify a group of nodes in the first subgraph in a chain structure;

wherein the instructions cause the first device to assign the group of nodes to one stream.

14. The system of claim 11 , wherein the assigning comprises:

analyzing the first subgraph to identify a first node in the first subgraph has a plurality of directed edges as outputs;

wherein the instructions cause the first device to assign, for each of the directed edges, a node to which the directed edge points to a unique stream of the hardware accelerator.

15. The system of claim 11 , wherein the instructions cause the first device to determine, for each node, a respective amount of memory resources in the hardware accelerator consumed by the operation represented by the node based on the directed edges to the node, wherein the assigning is based at least on the respective amount of memory resources.

16. The system of claim 11 , wherein the instructions cause the first device to determine a particular operation represented by a node has finished at a particular stream;

in response to determining the particular operation has finished:

determine a first amount of memory consumed by the particular operation that will be freed;

determine, for each of a group of unassigned nodes, a respective estimated amount of memory consumed by the unassigned node;

determine, from the group of unassigned nodes, a first unassigned node with the estimated amount of memory that maximizes usage of the first amount of memory; and

assign an operation represented by the first unassigned node to the particular stream.

17. A computer program product encoded on one or more non-transitory computer storage media, the computer program product comprising instructions that when executed by a hardware accelerator having a plurality of streams cause the hardware accelerator to perform operations comprising:

receiving, by a computational graph system, a request to perform a computation represented by a computational graph;

obtaining data representing multiple subgraphs of the computational graph, the computational graph comprising a plurality of nodes and directed edges, wherein each node represents a respective operation, wherein each directed edge connects a respective first node to a respective second node that represents an operation that receives, as input, an output of an operation represented by the respective first node;

assigning the multiple subgraphs to respective devices of a plurality of devices, including assigning a first subgraph to a first device by a placer in the computational graph system;

determining that the first device comprises a hardware accelerator having a plurality of streams, wherein each stream is an independent hardware queue that performs operations assigned to the stream in the order assigned;

in response to determining that the first device comprises a hardware accelerator having a plurality of streams, generating instructions that when executed by the first device cause the first device to:

assign the operation represented by each node in the first subgraph to a respective stream in the plurality of streams of the hardware accelerator, including identifying two operations that do not depend on each other and assigning the two operations to different streams, and assigning to a single stream a first operation that changes an internal state of the hardware in a way that must happen before a second operation executes and assigning the second operation to be performed in the single stream after the first operation is complete; and

perform the operations represented by the nodes in the first subgraph in accordance with the assignment;

providing the instructions and the data to the first device;

sending to the first device a request causing the first device to perform the operations represented by the nodes in the first subgraph;

obtaining from the first device outputs of operations performed by the first device;

causing the other devices of the plurality of devices to perform the operations of the subgraphs assigned to the other devices; and

responding to the request to perform the computation with one or more outputs obtained from the plurality of devices.

18. The computer program product of claim 17 , wherein the request specifies identifying one or more particular outputs from one or more respective nodes in the subgraph, further comprising:

receiving, from the first device, the one or more particular outputs; and

providing the one or more particular outputs to the client.

19. The computer program product of claim 17 , further comprising:

analyzing the first subgraph to identify a group of nodes in the first subgraph in a chain structure;

wherein the instructions cause the first device to assign the group of nodes to one stream.

20. The computer program product of claim 17 , wherein the assigning comprises:

analyzing the first subgraph to identify a first node in the first subgraph has a plurality of directed edges as outputs;

wherein the instructions cause the first device to assign, for each of the directed edges, a node to which the directed edge points to a unique stream of the hardware accelerator.

21. The computer program product of claim 17 , wherein the instructions cause the first device to determine, for each node, a respective amount of memory resources in the hardware accelerator consumed by the operation represented by the node based on the directed edges to the node, wherein the assigning is based at least on the respective amount of memory resources.

22. The computer program product of claim 17 , wherein the instructions cause the first device to determine a particular operation represented by a node has finished at a particular stream;

in response to determining the particular operation has finished:

determine a first amount of memory consumed by the particular operation that will be freed;

determine, for each of a group of unassigned nodes, a respective estimated amount of memory consumed by the unassigned node;

determine, from the group of unassigned nodes, a first unassigned node with the estimated amount of memory that maximizes usage of the first amount of memory; and

assign an operation represented by the first unassigned node to the particular stream.

23. A method for distributing a computation represented by a computational subgraph over a plurality of streams of a stream-based accelerator device and performing the computation including performing the computation represented by the subgraph on streams of the stream-based accelerator device, the method comprising:

receiving, by a hardware accelerator having a plurality of streams, data representing a subgraph of a computational graph, the computational graph comprising a plurality of nodes and directed edges, wherein each node represents a respective operation, wherein each directed edge connects a respective first node to a respective second node that represents an operation that receives, as input, an output of an operation represented by the respective first node, the subgraph assigned to a hardware accelerator by a placer in a computational graph system, wherein each stream of the plurality of streams is an independent hardware queue that performs operations assigned to the stream in the order assigned;

assigning, by the hardware accelerator, the operation represented by each node in the subgraph to a respective stream in the plurality of streams of the hardware accelerator, including identifying two operations that do not depend on each other and assigning the two operations to different streams, and assigning to a single stream a first operation that changes an internal state of the hardware in a way that must happen before a second operation executes and assigning the second operation to be performed in the single stream after the first operation is complete;

performing, by the hardware accelerator, the operations represented by the nodes in the subgraph in accordance with the assignment;

receiving a request identifying one or more particular outputs from one or more respective nodes in the subgraph; and

providing the one or more particular outputs to a client.

24. The method of claim 23 , further comprising:

receiving data identifying a group of nodes in the subgraph in a chain structure; and

assigning the group of nodes to one stream.

25. The method of claim 23 , wherein the assigning comprises:

receiving data identifying a first node in the subgraph having a plurality of directed edges as outputs; and

assigning, for each of the directed edges, a node to which the directed edge points to a unique stream of the hardware accelerator.

26. The method of claim 23 , further comprising determining, for each node, a respective amount of memory resources in the hardware accelerator consumed by the operation represented by the node based on the directed edges to the node, wherein the assigning is based at least on the respective amount of memory resources.

27. The method of claim 23 , further comprising determining a particular operation represented by a node has finished at a particular stream;

in response to determining the particular operation has finished, determining a first amount of memory consumed by the particular operation that will be freed;

determining, for each of a group of unassigned nodes, a respective estimated amount of memory consumed by the unassigned node;

determining, from the group of unassigned nodes, a first unassigned node with the estimated amount of memory that maximizes usage of the first amount of memory; and

assigning an operation represented by the first unassigned node to the particular stream.