Efficient execution of machine learning models in heterogeneous processing environments

Abstract

Certain aspects provide techniques and apparatuses for efficient operation of a machine learning model in a heterogeneous computing environment. An example method includes partitioning a graph representing a machine learning model into a plurality of subgraphs. Each subgraph generally represents a portion of the machine learning model. For each subgraph, a plurality of execution paths are simulated based on permutations of using different processing unit types to execute portions of the subgraph and starting with each input source processing unit type selected from the different processing unit types, and an execution path having a lowest cost is selected from the plurality of execution paths. The machine learning model is implemented based on the selected execution path for each subgraph.

What is claimed is: 
     
         1 . A processor-implemented method, comprising:
 partitioning a graph representing a machine learning model into a plurality of subgraphs, each subgraph representing a portion of the machine learning model;   for each subgraph, simulating a plurality of execution paths based on permutations of using different processing unit types to execute portions of the subgraph and starting with each input source processing unit type selected from the different processing unit types;   for each subgraph, selecting an execution path from the plurality of execution paths having a lowest cost; and   implementing the machine learning model based on the selected execution path for each subgraph.   
     
     
         2 . The method of  claim 1 , wherein partitioning the graph representing the machine learning model into the plurality of subgraphs comprises partitioning the graph representing the machine learning model into a first set of subgraphs associated with a first processing unit type of the different processing unit types and a second set of subgraphs associated with a second processing unit type of the different processing unit types. 
     
     
         3 . The method of  claim 2 , wherein the first set of subgraphs and the second set of subgraphs comprise graphs generated based on a common fusion boundary across a first processing system associated with the first processing unit type and a second processing system associated with the second processing unit type. 
     
     
         4 . The method of  claim 3 , wherein the common fusion boundary comprises a point in the machine learning model at which a subgraph in the first set of subgraphs and a corresponding subgraph in the second set of subgraphs output a common output for ingestion into a subsequent portion of the machine learning model. 
     
     
         5 . The method of  claim 1 , wherein simulating the plurality of execution paths for each subgraph comprises simulating an execution time for executing operations identified in the subgraph including context switching time for transitions from a first processing unit type of the different processing unit types to a second processing unit type of the different processing unit types. 
     
     
         6 . The method of  claim 1 , further comprising modifying a subgraph from the plurality of subgraphs based on combining consecutive portions of a subgraph representing operations for which execution should remain with the same processing unit type, wherein the plurality of execution paths are simulated based on the modified subgraph. 
     
     
         7 . The method of  claim 1 , wherein selecting the execution path for each subgraph comprises selecting the execution path having the lowest cost for each subgraph based on a backwards traversal of a graph representing the simulated plurality of execution paths for the plurality of subgraphs representing the machine learning model. 
     
     
         8 . The method of  claim 1 , further comprising generating an inference using the implemented machine learning model based on an input into the implemented machine learning model and the selected execution path for each subgraph. 
     
     
         9 . The method of  claim 1 , wherein the different processing unit types comprise two or more of a central processing unit (CPU), a graphics processing unit (GPU), or a neural processing unit (NPU). 
     
     
         10 . A processing system, comprising:
 at least one memory having executable instructions stored thereon; and   one or more processors configured to execute the executable instructions to cause the processing system to:
 partition a graph representing a machine learning model into a plurality of subgraphs, each subgraph representing a portion of the machine learning model; 
 for each subgraph, simulate a plurality of execution paths based on permutations of using different processing unit types to execute portions of the subgraph and starting with each input source processing unit type selected from the different processing unit types; 
 for each subgraph, select an execution path from the plurality of execution paths having a lowest cost; and 
 implement the machine learning model based on the selected execution path for each subgraph. 
   
     
     
         11 . The processing system of  claim 10 , wherein to partition the graph representing the machine learning model into the plurality of subgraphs, the one or more processors are configured to cause the processing system to partition the graph representing the machine learning model into a first set of subgraphs associated with a first processing unit type of the different processing unit types and a second set of subgraphs associated with a second processing unit type of the different processing unit types. 
     
     
         12 . The processing system of  claim 11 , wherein the first set of subgraphs and the second set of subgraphs comprise graphs generated based on a common fusion boundary across a first processing system associated with the first processing unit type and a second processing system associated with the second processing unit type. 
     
     
         13 . The processing system of  claim 12 , wherein the common fusion boundary comprises a point in the machine learning model at which a subgraph in the first set of subgraphs and a corresponding subgraph in the second set of subgraphs output a common output for ingestion into a subsequent portion of the machine learning model. 
     
     
         14 . The processing system of  claim 10 , wherein to simulate the plurality of execution paths for each subgraph, the one or more processors are configured to cause the processing system to simulate an execution time for executing operations identified in the subgraph including context switching time for transitions from a first processing unit type of the different processing unit types to a second processing unit type of the different processing unit types. 
     
     
         15 . The processing system of  claim 10 , wherein the one or more processors are further configured to cause the processing system to modify a subgraph from the plurality of subgraphs based on combining consecutive portions of a subgraph representing operations for which execution should remain with the same processing unit type, wherein the plurality of execution paths are simulated based on the modified subgraph. 
     
     
         16 . The processing system of  claim 10 , wherein to select the execution path for each subgraph, the one or more processors are configured to cause the processing system to select the execution path having the lowest cost for each subgraph based on a backwards traversal of a graph representing the simulated plurality of execution paths for the plurality of subgraphs representing the machine learning model. 
     
     
         17 . The processing system of  claim 10 , wherein the one or more processors are further configured to cause the processing system to generate an inference using the implemented machine learning model based on an input into the implemented machine learning model and the selected execution path for each subgraph. 
     
     
         18 . The processing system of  claim 10 , wherein the different processing unit types comprise two or more of a central processing unit (CPU), a graphics processing unit (GPU), or a neural processing unit (NPU). 
     
     
         19 . A processing system, comprising:
 means for partitioning a graph representing a machine learning model into a plurality of subgraphs, each subgraph representing a portion of the machine learning model;   means for simulating, for each subgraph, a plurality of execution paths based on permutations of using different processing unit types to execute portions of the subgraph and starting with each input source processing unit type selected from the different processing unit types;   means for selecting, for each subgraph, an execution path from the plurality of execution paths having a lowest cost; and   means for implementing the machine learning model based on the selected execution path for each subgraph.   
     
     
         20 . The processing system of  claim 19 , wherein the means for partitioning the graph representing the machine learning model into the plurality of subgraphs comprises means for partitioning the graph representing the machine learning model into a first set of subgraphs associated with a first processing unit type of the different processing unit types and a second set of subgraphs associated with a second processing unit type of the different processing unit types. 
     
     
         21 . The processing system of  claim 20 , wherein the first set of subgraphs and the second set of subgraphs comprise graphs generated based on a common fusion boundary across a first processing system associated with the first processing unit type and a second processing system associated with the second processing unit type. 
     
     
         22 . The processing system of  claim 21 , wherein the common fusion boundary comprises a point in the machine learning model at which a subgraph in the first set of subgraphs and a corresponding subgraph in the second set of subgraphs output a common output for ingestion into a subsequent portion of the machine learning model. 
     
     
         23 . The processing system of  claim 19 , wherein the means for simulating the plurality of execution paths for each subgraph comprises means for simulating an execution time for executing operations identified in the subgraph including context switching time for transitions from a first processing unit type of the different processing unit types to a second processing unit type of the different processing unit types. 
     
     
         24 . The processing system of  claim 19 , further comprising means for modifying a subgraph from the plurality of subgraphs based on combining consecutive portions of a subgraph representing operations for which execution should remain with the same processing unit type, wherein the plurality of execution paths are simulated based on the modified subgraph. 
     
     
         25 . The processing system of  claim 19 , wherein the means for selecting the execution path for each subgraph comprises means for selecting the execution path having the lowest cost for each subgraph based on a backwards traversal of a graph representing the simulated plurality of execution paths for the plurality of subgraphs representing the machine learning model. 
     
     
         26 . The processing system of  claim 19 , further comprising means for generating an inference using the implemented machine learning model based on an input into the implemented machine learning model and the selected execution path for each subgraph. 
     
     
         27 . The processing system of  claim 19 , wherein the different processing unit types comprise two or more of a central processing unit (CPU), a graphics processing unit (GPU), or a neural processing unit (NPU). 
     
     
         28 . A non-transitory computer-readable medium having executable instructions stored thereon which, when executed by one or more processors, perform an operation comprising:
 partitioning a graph representing a machine learning model into a plurality of subgraphs, each subgraph representing a portion of the machine learning model;   for each subgraph, simulating a plurality of execution paths based on permutations of using different processing unit types to execute portions of the subgraph and starting with each input source processing unit type selected from the different processing unit types;   for each subgraph, selecting an execution path from the plurality of execution paths having a lowest cost; and   implementing the machine learning model based on the selected execution path for each subgraph.