Dynamic pathways for artificial intelligence and tensor computation graphs
A method for optimizing control flow in compiled computation graphs includes defining an intermediate representation (IR) of a computation graph, the computation graph IR including a main computation graph having at least one control flow primitive layer node pointing to one or more control flow sub-graph nodes. Fusable layer nodes of the main computation graph are identified and removed from the main computation graph, and the removed fusable layer nodes are duplicated into each of the one or more control flow sub-graph nodes. The method can be applied to machine learning (ML) and artificial intelligence (AI) frameworks, for example, to support decision making and optimization of targets and scientific computations in a variety of applications such as in medical/healthcare (e.g., AI assisted healthcare or drug development).
1 . A computer-implemented for optimizing control flow in compiled computation graphs, the method comprising:
defining an intermediate representation (IR) of a computation graph, the computation graph IR including a main computation graph having at least one control flow primitive layer node pointing to one or more control flow conditional-graph nodes;
identifying fusable layer nodes of the main computation graph; and
removing the fusable layer nodes from the main computation graph and duplicating the removed fusable layer nodes into each of the one or more control flow conditional-graph nodes.
2 . The method of claim 1 , wherein layer nodes of the computation graph IR define code paths, and wherein the method further comprises, during runtime:
detecting one or more code paths that depend on one or more parameters that may change; and
enabling or disabling executing the one or more code paths based on the one or more parameters having changed during runtime.
3 . The method of claim 1 , further comprising copying to a memory a computation result of a layer node of the main computation graph executed prior to the at least one control flow primitive layer node.
4 . The method of claim 1 , wherein only a layer node that receives data or calls for data is configured to allocate or free memory for the data and for variables.
5 . The method of claim 4 , wherein each of the one or more control flow conditional-graph nodes is allowed to only allocate memory for itself.
6 . The method of claim 1 , further comprising adding a termination path sub-graph node that provides for all model outputs of the termination path sub-graph to be allocated and/or zeroed out, wherein the termination path sub-graph is executed at runtime in response to a zero-element tensor.
7 . The method of claim 1 , wherein the fusible layer nodes comprise operations that occur before and/or after execution of the at least one control flow primitive layer node.
8 . The method of claim 1 , wherein the defining the IR of the computation graph includes receiving the computation graph and processing the computation graph to produce the IR of the computation graph.
9 . The method of claim 1 , wherein the at least one control flow primitive layer node comprises a conditional construct.
10 . The method of claim 9 , wherein the conditional construct comprises a While construct or an If construct or a loop construct.
11 . The method of claim 9 , wherein the one or more control flow conditional-graph nodes comprises a first conditional-graph node and a second conditional-graph node, wherein the conditional construct is associated with the first conditional-graph node and the second conditional-graph node, and wherein only one of the first conditional-graph node or the second conditional-graph node is executed during runtime based on a dynamic runtime condition.
12 . The method of claim 11 , wherein the conditional construct is an If construct, and wherein the first conditional-graph node is associated with a true condition of the If construct and the second conditional-graph node is associated with a false condition of the if construct.
13 . The method of claim 1 , wherein the main computation graph comprises the at least one control flow primitive layer node and a plurality of non-control flow primitive layer nodes, and wherein each of the identified fusable layer nodes is a non-control flow primitive layer node from the plurality of non-control flow primitive layer nodes.
14 . The method of claim 13 , wherein identifying the fusable layer nodes of the main computation graph is based on analyzing the plurality of non-control flow primitive layer nodes that occur before or after each of the at least one control flow primitive layer node.
15 . A system for optimizing control flow in compiled computation graphs, the system comprising one or more hardware processors which, alone or in combination, are configured to provide for execution of the following steps:
defining an intermediate representation (IR) of a computation graph, the computation graph IR including a main computation graph having at least one control flow primitive layer node pointing to one or more control flow conditional-graph nodes;
identifying fusable layer nodes of the main computation graph; and
removing the fusable layer nodes from the main computation graph and duplicating the removed fusable layer nodes into each of the one or more control flow conditional-graph nodes.
16 . The system of claim 15 , wherein the steps further include adding a termination path sub-graph node that provides for all model outputs of the termination path sub-graph to be allocated and/or zeroed out, wherein the termination path sub-graph is executed at runtime in response to a zero-element tensor.
17 . The system of claim 15 , wherein the steps further include copying to a memory a computation result of a layer node of the main computation graph executed prior to the at least one control flow primitive layer node.
18 . The system of claim 15 , wherein layer nodes of the computation graph IR define code paths, and wherein the steps further include, during runtime:
detecting one or more code paths that depend on one or more parameters that may change; and
enabling or disabling executing the one or more code paths based on the one or more parameters having changed during runtime.
19 . A tangible, non-transitory computer-readable medium having instructions thereon which, upon being executed by one or more hardware processors, alone or in combination, provide for execution of the following steps:
defining an intermediate representation (IR) of a computation graph, the computation graph IR including a main computation graph having at least one control flow primitive layer node pointing to one or more control flow conditional-graph nodes;
identifying fusable layer nodes of the main computation graph; and
removing the fusable layer nodes from the main computation graph and duplicating the removed fusable layer nodes into each of the one or more control flow conditional-graph nodes.