Apparatus and method for processing a computational graph representing task execution scheduling

Abstract

An apparatus (100) for processing a computation graph executed on a device (150), the device comprising a plurality of processing elements (151, 152, 153, 154) having access to a memory (155), the apparatus comprising a processor (101) configured to: receive (1001) information about the elements, each element being associated with a type of task; receive (1002) information indicating a size of the memory; receive (1003) a graph having nodes corresponding to executable tasks corresponding to a type and constrained to be executed in an element associated with its type, and arcs indicating precedence constraints; receive (1004) original task total orders, each corresponding to one element and respecting the precedence constraints; determine (1005) new task total orders, each corresponding to one element and respecting the precedence constraints, such that a peak memory usage of the execution of the graph does not exceed the size of the memory. This can make minimal adjustments to a given task ordering, thereby avoiding memory shortage problems.

Description

Technical Field

[0001] This invention relates to the processing of a computation graph representing the execution schedule of multiple tasks, for example, to determine a task execution schedule that is memory efficient and requires the fewest reordering operations. Background Technology

[0002] In recent years, the size of neural network models has been steadily and rapidly increasing. State-of-the-art Transformer-based models (such as the paper by Vaswani et al. published in Volume 30 of *Advances in Neural Information Processing Systems* in 2017) have seen significant growth. The model described in "Attention Is All You Need" now has nearly or even more than one trillion parameters.

[0003] As models grow larger, more memory is needed to allocate space for parameters and intermediate data (tensors) generated and consumed by neural operators during training, especially when training uses large batch sizes. While the available memory resources in state-of-the-art artificial intelligence (AI) accelerators are constantly increasing, their speed cannot keep pace with the expansion of model size.

[0004] The model size has increased approximately 410-fold every two years, while memory resources have roughly doubled during the same period. Therefore, more and more devices and corresponding memory resources are needed to train such large models, leading to increased design complexity and operating costs.

[0005] To overcome this memory barrier and achieve efficient execution, various alternative methods have been developed, such as optimizing parallel design, memory offloading, model pruning and quantization, intermediate data recomputation, memory-aware task scheduling, and optimizing tensor memory allocation.

[0006] Nevertheless, optimizing for memory can incur computational overhead and negatively impact training performance. Exploring the trade-off between improving training throughput and saving memory is one of the biggest challenges in designing efficient deep neural network (DNN) execution systems.

[0007] We hope to develop a method that can at least partially solve the above problems. Summary of the Invention

[0008] According to one aspect, an apparatus is provided for processing a computation graph executed on a device, the device including a plurality of processing elements, each processing element having access to shared memory, the apparatus including one or more processors, the processors being configured to: receive information about the plurality of processing elements of the device, each processing element being associated with a corresponding task type among a plurality of task types; receive information indicating the size of the shared memory of the device; receive a computation graph having a plurality of nodes corresponding to executable tasks and arcs indicating priority constraints between the tasks, each task corresponding to a corresponding task type among the plurality of task types and constrained to be executed in the processing element associated with its corresponding task type; receive a plurality of original task full sequences, each original task full sequence corresponding to a corresponding processing element, wherein the plurality of original task full sequences comply with the priority constraints; determine a plurality of new task full sequences such that the peak memory usage for executing the computation graph does not exceed the size of the shared memory, each new task full sequence corresponding to a corresponding processing element, wherein the plurality of new task full sequences comply with the priority constraints.

[0009] This allows for minimal modifications to the task execution scheduling of a set of processing elements, thus adhering to shared memory constraints during parallel execution. Scheduling efficiency can be maintained while ensuring scheduling feasibility in terms of memory requirements. This method also enables optimization of task ordering without starting from scratch. Therefore, this approach allows for minimal adjustments to the given task ordering of the computation graph, helping to avoid memory shortages during computation graph execution.

[0010] The one or more processors can be used to: convert the plurality of new tasks into a total order of operations (TOE) multiple non-redundant task dependencies; and insert the multiple non-redundant task dependencies into the computation graph. This allows for the elimination of redundant task dependencies, thereby improving the execution efficiency of the computation graph.

[0011] The one or more processors can be used to: compile the computation graph; and send the compiled computation graph to the device for execution via the plurality of processing elements and the shared memory. This allows the compiled computation graph to be executed within the processing elements of the device, without exceeding the capacity of the shared memory.

[0012] The one or more processors can be used to convert the plurality of new tasks sequentially into a plurality of non-redundant task dependencies in the following manner: for each processing element, when traversing the corresponding total order of the new tasks for the respective processing element, for each pair of consecutive tasks, if there is no directed path of computation graph arcs from the previous task to the next task in the corresponding pair of consecutive tasks, then a task dependency is defined. This enables the removal of redundant task dependencies.

[0013] The one or more processors can be used to insert the multiple non-redundant task dependencies into the computation graph by introducing a new arc to the computation graph for each dependency. This enables the computation graph to be compiled and executed in total order based on multiple new tasks.

[0014] The one or more processors can be used to determine the order of the plurality of new tasks by: estimating the computation time and memory consumption of each task; estimating the memory requirements during the execution of the computation graph based on the total order of the plurality of original tasks; and determining the total order of the plurality of new tasks based on the estimated computation time and memory consumption of each task and the estimated memory requirements during the execution of the computation graph. This enables the computation graph to be executed in a computationally efficient and memory-efficient manner.

[0015] For each processing element, the corresponding original full sequence of tasks can correspond to the execution order of the plurality of tasks associated with that processing element. This allows for optimization of the task ordering without having to start from scratch.

[0016] The total sequence of the multiple new tasks can have minimal deviation from the total sequence of the multiple original tasks in a certain metric. This metric can be the difference or relative error between the estimated execution times of the corresponding computational graphs for the total sequences of the multiple original tasks and the total sequences of the multiple new tasks. The metric can be Kendall Jenner's algorithm. Rank distance.

[0017] The computation graph can represent a neural network. The input computation graph can be generated from the end user's model file (e.g., a neural network model file).

[0018] The computation graph can be a directed acyclic graph. In this case, the arcs of the graph are directed arcs (pointing from the first node to the next node). This allows for the definition of the task execution order.

[0019] The computation graph can be a hierarchical graph. The hierarchical graph can include subgraphs, each subgraph corresponding to an executable task.

[0020] The one or more processors can be used to make minimal modifications to the task execution schedule of multiple processing elements, thereby adhering to device shared memory constraints. The one or more processors can be used to output a fully ordered sequence of tasks with minimal modifications to each processing element and in compliance with constraints, so that the compiled computation graph can be executed through the shared memory. This may help avoid out-of-memory issues during the execution of the computation graph on the device.

[0021] The multiple raw task sequences (one for each processing element) can indicate user-defined scheduling decisions or automatic scheduling decisions. These raw task sequences can then be optimized for memory consumption.

[0022] At least one of the multiple types can be one of tensor computation, vector element computation, cubic element computation, and communication. Each processing element can be associated with a type to perform a task of that type.

[0023] The one or more processors can be used to form the plurality of new full orders of tasks to maintain the efficiency (e.g., completion time efficiency) of the original schedule defined by the plurality of original full orders of tasks. This allows the device to achieve efficiency optimization within shared memory constraints when executing the computation graph.

[0024] The computation graph can represent the execution schedule of the multiple tasks. The execution schedule can be a typed parallel execution schedule.

[0025] The total order of the multiple original tasks can be determined by performing a compilation optimization process based on the received computation graph, received processing elements, and shared memory information. This enables the definition of the original execution schedule.

[0026] Executing the compiled computation graph on the device allows for the training or inference of deep neural networks. This enables efficient training or inference to be performed on the device.

[0027] According to another aspect, a computer-implemented method is provided for processing a computation graph executed on a device, the device including a plurality of processing elements, each processing element having access to shared memory, the method comprising: receiving information about the plurality of processing elements of the device, each processing element being associated with a corresponding task type among a plurality of task types; receiving information indicating the size of the shared memory of the device; receiving a computation graph having a plurality of nodes corresponding to executable tasks and arcs indicating priority constraints between the tasks, each task corresponding to a corresponding task type among the plurality of task types and constrained to execution in the processing element associated with its corresponding task type; receiving a plurality of original task full sequences, each original task full sequence corresponding to a corresponding processing element, wherein the plurality of original task full sequences comply with the priority constraints; determining a plurality of new task full sequences such that the peak memory usage for executing the computation graph does not exceed the size of the shared memory, each new task full sequence corresponding to a corresponding processing element, wherein the plurality of new task full sequences comply with the priority constraints.

[0028] This allows for minimal modifications to the task execution scheduling of a set of processing elements, thus adhering to shared memory constraints during parallel execution. Scheduling efficiency can be maintained while ensuring scheduling feasibility in terms of memory requirements. This method also enables optimization of task ordering without starting from scratch. Therefore, this approach allows for minimal adjustments to the given task ordering of the computation graph, helping to avoid memory shortages during computation graph execution.

[0029] According to another aspect, a computer program is provided that, when executed by a computing device, causes the computing device to perform the above-described method.

[0030] According to another aspect, a data carrier for storing the aforementioned computer program in a non-transient form is provided. Attached Figure Description

[0031] The invention will now be described by way of example with reference to the accompanying drawings.

[0032] In the attached diagram:

[0033] Figure 1 An example of a computer system including a host and devices is illustrated schematically;

[0034] Figures 2(a) and 2(b) schematically illustrate examples of computational graphs;

[0035] Figure 3 A device with shared memory and multiple processing elements is schematically shown;

[0036] Figure 4 The steps of an exemplary method for processing a computation graph and executing the graph on a device are illustrated schematically;

[0037] Figure 5 The steps of an exemplary method for determining the total order of a new task that complies with constraints are illustrated schematically.

[0038] Figure 6(a) schematically illustrates an example of the primitive execution schedule of the total order of multiple primitive tasks indicating the tasks of the computation graph;

[0039] Figure 6(b) schematically illustrates the shared memory usage and efficiency corresponding to the execution of a computation graph with a total order of multiple original tasks as shown in Figure 6(a);

[0040] Figure 7(a) schematically illustrates an example of task execution scheduling using a computation graph with multiple new task total-order modifications as defined herein;

[0041] Figure 7(b) schematically illustrates the shared memory usage and efficiency corresponding to the execution of a computation graph with a full order of multiple new tasks as shown in Figure 7(a);

[0042] Figure 8 The steps of an exemplary method for transforming a computation graph by inserting non-redundant task dependencies are illustrated schematically.

[0043] Figures 9(a) and 9(b) illustrate how the total order of multiple new tasks shown in Figure 7(a) is converted into a non-redundant set of task dependencies and inserted into the computation graph;

[0044] Figure 10 The steps of a computational graph processing method according to some embodiments of the present invention are shown. Detailed Implementation

[0045] Embodiments of the present invention relate to the processing of computational graphs.

[0046] This paper describes a method for modifying task execution scheduling so that the task execution schedule can be adapted to the available memory resources of the device on which the task is executed.

[0047] In some examples, task execution scheduling can be a parallel execution schedule for DNN training. The goal is to determine the minimum rescheduling of tasks executed on devices (e.g., heterogeneous AI accelerators) under task type constraints. Such optimized rearrangements can save sufficient memory without negatively impacting computation.

[0048] Traditionally, optimizing task scheduling can improve the execution performance of DNNs. However, this can lead to increased memory requirements. Suppose a task scheduling scheme optimized for DNN training has been pre-computed or manually defined. It may be desirable to execute this scheduling on a given piece of hardware with limited shared memory. It is desirable to avoid recompiling the schedule from scratch, for example, to avoid computational overhead or to preserve the original scheduling decisions as much as possible. It is desirable to make minimal modifications to the given schedule to fit the available memory, that is, to avoid significant changes to prevent a significant performance degradation during execution. This problem will be formally described below in a heterogeneous (typed) environment.

[0049] Exemplary system architecture such as Figure 1 As shown. The system includes a host 100 and a device 150. The host 100 is an apparatus for processing graphs to determine the full sequence of new tasks. The host includes a processing unit 101, a memory 102, and an optional disk storage 103. The memory 102 and / or the disk storage 103 may store code executable by the processor 101 in a non-transitory manner to implement the methods described herein. In some implementations, the host 100 and the device 150 may each include a transceiver capable of communicating with the other of the host 100 and the device 150 via a network. This network may be a publicly accessible network such as the Internet. Alternatively, the network may be a local area network (LAN).

[0050] The methods described in this article can be deployed in various ways, such as in the cloud, on a host computing device 100, or in dedicated hardware.

[0051] Device 150 includes multiple PEs 151, 152, 153, and 154. Each PE is associated with a different task type. Device 150 also includes shared memory 155. Shared memory 155 can have a predetermined size. Typical PE types include cube (e.g., Ascend910 DaVinci Core), vector, communication, or other specialized types. Cube and vector PEs can be referred to as computation PEs. All PEs 151, 152, 153, and 154 can access shared memory 155.

[0052] Consider a computational graph (CG) as input, in this example, where the computational graph models the DNN. In this example, the CG is a Directed Acyclic Graph (DAG) that includes nodes corresponding to neural operators (in other words, performable tasks). The graph also includes arcs between nodes, which define priority constraints between tasks. Each task is associated with one of several types, such as tensor computation, vector unit computation, cubic unit computation, and communication (e.g., aggregation and reduction).

[0053] Input CG, device information (i.e., information indicating the size of the shared memory and the task type associated with each PE), and multiple raw task full sequences can be stored in disk storage 103 or loaded into storage 102.

[0054] The output CG enhanced by the new mission can be determined by the processing unit 101 of the host device 100.

[0055] CG is designed to execute on a device 150 comprising multiple PEs 151, 152, 153, 154 and shared memory 155 (which may be a single shared memory). The shared memory can be logical or physical memory. The size of the shared memory may be limited. Each PE is associated with a specific task type. Each task is restricted to execution within a unique PE that matches its type. Device information can be provided as input, indicating the size of the shared memory and the task type associated with each PE.

[0056] For a CG, there are multiple primitive task total orders (one for each PE), all of which adhere to all priority constraints. Each task total order is associated with multiple tasks. For each PE, the task total order corresponds to the execution order of its associated multiple tasks (that is, multiple tasks matching the PE type).

[0057] Given the above input, the goal is to output multiple new task total orders (one for each PE), which adhere to priority constraints such that the peak memory usage of CG execution does not exceed the size of the device's shared memory (which may be finite). It is also desirable that the multiple task total orders correlate with the multiple original task total orders under certain appropriate metrics (e.g., Kendall's algorithm). There is a minimum deviation in terms of rank distance.

[0058] Input CG can be generated from the end-user's neural network model file. For example, input CG can be generated from a Python file that defines a natural language processing (NLP) neural network model, or from an upstream parallel module (e.g., the MindSpore AutoPar optimizer) running on a cluster of devices, which can determine the CG for each device.

[0059] The input CG can consist of multiple operators of different PE types. The input CG can have arbitrary topology. In some cases, each node of the graph represents a corresponding executable task. In other cases, the CG can be a hierarchical graph, where some executable tasks correspond to multiple branching subgraphs. Each branch can be of any form. In some implementations, the upstream parallel module can force the branches to be symmetric.

[0060] Figures 2(a) and 2(b) illustrate examples of general input CGs 200 and 250 with parallel branches, respectively. Each CG 200, 250 includes multiple nodes. For these nodes, different shapes (ellipse, rectangle, rhombus, and hexagon) represent different PE types of nodes / tasks assigned for execution.

[0061] Each task corresponding to a node corresponds to a specific task type and is restricted to execution within the PE associated with its corresponding task type.

[0062] For example, in Figure 2(a), CG 200 includes nodes 201-212. The executable tasks corresponding to nodes 201, 202, 205, and 212 are executed on PEs with a first task type (represented by an ellipse in this example). The executable tasks corresponding to nodes 203 and 211 are executed on PEs with a second task type (represented by a rectangle in this example). The executable tasks corresponding to nodes 204, 206, 207, and 208 are executed on PEs with a third task type (represented by a rhombus in this example). The executable tasks corresponding to nodes 209 and 210 are executed on PEs with a fourth task type (represented by a hexagon in this example).

[0063] For example, in Figure 2(b), CG 250 includes nodes 251-264. The executable tasks corresponding to nodes 251, 252, 253, 254, and 264 are executed on PEs with a first task type. The executable tasks corresponding to nodes 261, 262, and 263 are executed on PEs with a second task type. The executable tasks corresponding to nodes 255, 256, and 257 are executed on PEs with a third task type. The executable tasks corresponding to nodes 258, 259, and 260 are executed on PEs with a fourth task type.

[0064] like Figure 3 As shown, in device 300, each PE 301, 302, 303, 304 associated with the corresponding task type can access shared memory 305.

[0065] Multiple raw task complete sequences (one for each PE) can be obtained through the compilation and optimization process. This process can take CG and device information as input. Alternatively, multiple raw task complete sequences for each PE can be manually inserted by the user.

[0066] In general software platforms for DNN execution (such as MindSpore, PyTorch, and TensorFlow), user-defined NN models undergo many processing steps before finally being executed on the target device.

[0067] Figure 4 The document outlines typical steps, illustrating the entire input parsing process chain 400 and other steps of the method.

[0068] In step 401, the input is parsed. The input includes:

[0069] (i) CG, for example, represents the user's NN model, where the nodes (tasks) of the CG represent neural operators, each neural operator is assigned to a specific PE type, and the arcs of the CG represent the data flow or control flow between two operators and indicate the priority constraints to be executed;

[0070] (ii) Device information, including multiple available PEs and their types and shared memory size;

[0071] (iii) Multiple tasks in total order (one for each PE), subject to CG priority constraints.

[0072] In mathematical symbolic representation, let... Indicates partial order, Represents the total order, where yes A subset of. If for any ,(in ), or If it is established, then it can be said for efficient.

[0073] Parsing the input CG, the partial order can be defined as follows: ,in It is a set of nodes (tasks) in CG, for any If there exists a sequence of length at least 1 from arrive A directed path (achieved by traversing the arcs in the input CG) is then... Established.

[0074] The following can be parsed: multiple available PEs, the size of shared memory in the device, and the processing type function that assigns each task in the CG to a PE. Each PE in the equipment (i.e. Total order of ) ,in It is all tasks ( (established) collection, for efficient.

[0075] In step 402, a total sequence of multiple new tasks that adheres to priority constraints is determined. In other words, a total sequence of multiple new tasks is determined based on priority constraints. The processing unit 101 can perform static analysis on the input to determine a total sequence of multiple new tasks (one for each PE), which still adheres to all priority constraints.

[0076] For each PE in the device The new total order can be determined as for efficient.

[0077] In step 403, each new task total order in the multiple new task total orders is converted into multiple non-redundant task dependencies. Then, the multiple non-redundant task dependencies are inserted into the CG.

[0078] In step 404, the modified CG is compiled in processing unit 101, thereby generating an executable file to be sent to devices 150 and 300.

[0079] In step 405, the executable CG is performed on device 150. The CG can be performed to, for example, perform DNN training or inference.

[0080] exist Figure 5 middle, Figure 4 Step 402 is broken down into a method including exemplary sub-steps 501-503.

[0081] In step 501, the computational resources required to execute each task are estimated (i) and (ii) the memory needed to store its output, as well as any possible workspace memory requirements. In some implementations, the computation time and memory consumption for each task can be estimated using a cost function or hardware performance analysis. Two natural numbers can be assigned to each task: the first number represents the estimated computation time, and the second number represents the memory expected to be allocated in shared memory. For example, these numbers can be calculated using a theoretical cost function or hardware performance analysis.

[0082] In step 502, the parallel execution implied by the input given in step 401 is analyzed. Using task computation intensity and memory estimates, the total memory required at each time point during CG execution can be estimated. Peak memory consumption time points can be identified. Using the above notation, the total sequence of inputs can be used... Estimate the memory requirements during CG execution.

[0083] In step 503, multiple new task full sequences (one for each PE) are determined such that their corresponding estimated peak memory consumption does not exceed the device's shared memory size. This minimizes the impact on the multiple original task full sequences. This can be done for each... Determine the new total order, for Effective. This can be achieved through mathematical programming or algorithm design, reducing peak memory usage to an executable level while maintaining... and The distance between them is the smallest.

[0084] This process can be optimized so that the number of task transpositions is close to the minimum feasible value compared to the original total task order. This problem can be described as a constraint programming (CP) or integer linear programming (ILP) model. Alternatively, algorithms can be designed to derive multiple new total task orders.

[0085] Figures 6(a), 6(b), 7(a) and 7(b) illustrate an application example of step 402, namely, determining the full sequence of multiple new tasks that comply with priority constraints.

[0086] Figure 6(a) illustrates the total sequence of multiple original tasks in the exemplary CG 600. The graph includes nodes 601-614, labeled S, 1-12, and E, respectively. CG 600 is a directed acyclic graph. Each arc is labeled with an integer representing the size of the output data (if any) of the source node. For simplicity, in this example, it is assumed that the integer corresponding to a node represents both the memory allocated to the corresponding task and its computational intensity.

[0087] In Figure 6(a), the total task order is shown as a sequence of dashed / dotted arcs. There are three original total task orders, including the tasks in the following order:

[0088] (a) (The additional arc is shown as a dashed arc)

[0089] (b) (The additional arcs are shown as dotted arcs)

[0090] (c) (The additional arc is shown as a dashed arc)

[0091] The task in task sequence (a) belongs to the first task type and is executed on one or more processing elements associated with the first task type (represented by an ellipse in this example). The task in task sequence (b) belongs to the second task type and is executed on one or more processing elements associated with the second task type (represented by a hexagon in this example). The task in task sequence (c) belongs to the third task type and is executed on one or more processing elements associated with the third task type (represented by a diamond in this example).

[0092] As shown in Figure 6(b), this graph illustrates the relationship between memory usage and completion time efficiency. Each task 1-12 consumes computation time and memory as indicated by boxes 651-662. The vertical axis corresponds to computation time, while the horizontal axis corresponds to memory usage. For the case where multiple original tasks are executed in total sequence to complete the computation graph, the estimated peak memory usage is 70, while the completion time efficiency is 160.

[0093] Boxes 671 and 672 consider the issue from a memory perspective rather than computation time, to cover the scenario where the memory associated with a task must be retained for a period of time after the task execution ends. In this example, for simplicity, a simple dataflow model is used, where output data resides in memory from the start of the production task until the latest start time of its consumption task. While simplified, such assumptions do not affect the generality of the approach, as described in the 2019 paper "Limiting thememory footprint when dynamically scheduling DAGs on shared-memory platforms" by Marchal, Simon, and Vivien, published in the *Journal of Parallel and Distributed Computing*, Volume 128, pp. 30–42, 2019, this model can simulate other common and more complex models.

[0094] Figures 7(a) and 7(b) show the complete sequence of several new missions of CG 600.

[0095] Figure 7(a) shows three new task total sequences, including the tasks in the following order:

[0096] (a)

[0097] (b)

[0098] (c)

[0099] As shown in Figure 7(b), this graph illustrates the relationship between memory usage and completion time efficiency. Each task 1-12 consumes computation time and memory as indicated by boxes 651-662. The vertical axis corresponds to computation time, while the horizontal axis corresponds to memory usage. As shown in Figure 7(b), for the case where multiple new tasks are executed in a total sequence to complete the computation graph, the estimated peak memory usage is 50, while the completion time efficiency is 190.

[0100] In the total sequence of tasks (a) and (c), tasks 5 and 11, as well as tasks 6 and 12, are transposed in their corresponding total sequence of operations (PEs) to accommodate parallel execution (the optimal solution) within a shared memory space of size 50. Similar to the descriptions of boxes 671 and 672 above, boxes 673 and 674 consider the issue from a memory perspective rather than a time perspective.

[0101] Figure 8 An exemplary application of step 403 is shown, namely, converting multiple new tasks into multiple non-redundant task dependencies and inserting them into CG.

[0102] In step 801, each task total order in the multiple task total orders is considered. If the CG priority constraint does not yet cover the dependency, a task dependency is defined for every two consecutive tasks in the total order.

[0103] For each PE In traversing the total order of its corresponding tasks At that time, for each pair of consecutive tasks ,if Then the definition starts from arrive Task dependencies.

[0104] In step 802, for example, by introducing one or more new arcs into the CG data structure, the determined non-redundant task dependencies are inserted into the CG. Each new arc corresponds to a determined task dependency. One new arc can be introduced for each dependency.

[0105] Figures 9(a) and 9(b) illustrate how a total sequence of multiple new tasks is converted into a non-redundant set of task dependencies and inserted into the CG.

[0106] For multiple new task total sequences as shown in Figure 7(a), in Figure 9(b), the CG only retains non-redundant arcs. That is, there is no directed path from the source node to the target node of the arc beforehand. Redundant arcs will be removed from the CG.

[0107] In Figure 9(b), new arcs 901, 902, 903, 904, and 905 are introduced into the CG data structure. In this example, new arcs 901, 902, 903, 904, and 905 are directed arcs.

[0108] An exemplary embodiment for determining a total order of multiple new tasks that comply with constraints will now be described. This embodiment employs CP to illustrate the optimization problem, as defined above.

[0109] In this example, the model's variables can be defined as follows:

[0110] - The term "occupy" refers to a slice of computation time that a task takes up during the execution of a PE.

[0111] - For tasks in CG ,set up Indicates the existence of from The set of tasks pointed to by the starting arc. Similarly, let... This indicates that a pointer exists. The task set of the arc.

[0112] - set up Representation and Task Related PE types, Indicates task The estimated calculation time.

[0113] - Due to priority constraints and Implicit priority constraints Therefore, the priority constraints between tasks are simplified to checking only the parent-child relationship.

[0114] - Assuming each task Ultimately, they will all be allocated a start time (i.e., ), and will be assigned to And in time slice (That is to say, it will not be preempted before its entire computation time.)

[0115] - For PE type This type of PE is represented as .

[0116] - For each task ,set up This represents the estimated memory allocated by the task to store its output.

[0117] - set up This represents the upper limit of the expected completion time of the new schedule. This constant can be estimated based on the completion time of the original schedule, or it can be determined using methods such as binary search.

[0118] - set up This represents a constant that matches the size of the shared memory.

[0119] - For each ,set up Indicates and The corresponding total order of input, This indicates the new total order to be determined. Add an index parameter to return the task at that specific location, for example, for the task in Figure 6(a) / Figure 7(a) with the first PE (i.e., The corresponding total order, True, assuming sorting starts from the index. In the beginning, Established.

[0120] Suppose that between two tasks in the same PE, there is a total sequence... This represents an index used to describe how much these two total orders differ from each other, for example, Kendall's... Rank distance (as described in Kendall MG's paper "A new measure of rank correlation" published in Biometrics, Vol. 30, pp. 81-93, 1938).

[0121] In this specific example, the constraint set is as follows:

[0122] (1) Each task is occupied on the PE of its associated type.

[0123] (2) A task only begins to occupy computation time slices after its parent task has been completed.

[0124] (3) The occupancy of any two tasks on the same PE cannot overlap.

[0125] (4) The estimated peak memory size must not exceed the size of the device’s shared memory.

[0126] (5) Assume that the task is greedily started on the device including the PE. Once the parent task of a task and the previous task in the total order of the PE (if it exists) have been executed, the task will start executing immediately.

[0127] Other constraints can be used in alternative implementations.

[0128] An example pseudocode is shown below.

[0129] Constraint procedure begins

[0130] Target

[0131]

[0132] constraint

[0133] :

[0134]

[0135]

[0136]

[0137]

[0138] (1)

[0139] (2) :

[0140]

[0141] (3) :

[0142]

[0143] (4)

[0144] (5) For each :

[0145]

[0146] For each integer :

[0147] ( )

[0148] End of constraint procedure

[0149] Consider applying the CP model to the examples shown in Figures 6(a), 6(b), 7(a), and 7(b). For demonstration purposes, in this example, for each task... , Established, among which Indicates task The size set of the output tensors.

[0150] In this example:

[0151] ,

[0152] ,

[0153] ,

[0154] ,

[0155] ,

[0156] .

[0157] Suppose there are three PEs, each of which is a different type, i.e. The complete sequence of input (original) and output (new) tasks is as follows:

[0158] -

[0159] -

[0160] -

[0161] -

[0162] -

[0163] -

[0164] The program even Even under these conditions, a solution can be obtained, and the optimization objective is reduced to a minimum. This is related to the output solution (Kendall). The two task transpositions in the rank distance correspond to: one in the total task order (a) ( ), once in the total sequence of tasks (c) ).

[0165] The above constraint procedure can be converted into an ILP by linearizing each nonlinear constraint and objective function. For example, constraint (3) can be linearized using the industry-standard "Big M" method commonly used for OR constraints in the literature. Constraint (4) can be linearized by considering only the ordering of "task start / end" events, for example, by the implication of the corresponding start / end variables, which can be represented as OR constraints and linearized accordingly. In this way, a lifecycle, i.e., "data start / end" events, can be defined, where data blocks are in a "live" state (binary variables) between their start and end events, and the constraint can be expressed as "for each start / end task event, the total size of the data that survives in that event must be less than or equal to the memory limit". Constraint (5) can be linearized as follows: Maintain an "order" variable for each task, the value of which is the task's position in the new order to be determined. Then, write down several implicit constraints with the following conditions: "For each task, if the order equals 1" (top part of the constraint); "For each pair of tasks, if their order is greater than 1, and the order of the preceding task is the order of the following task plus 1" (bottom part of the constraint). Constraints containing "max" can be linearized using standard techniques. If the objective function is not yet linearized, it can be linearized using techniques similar to those described above.

[0166] Another embodiment will now be described for determining a new total order of tasks that comply with constraints in the form of an algorithmic process. This method “traverses” the input schedule from beginning to end, identifying all bottlenecks—the portions of the schedule that require more memory than the limits. Once a bottleneck is identified, it is resolved by shifting the task causing the bottleneck after the next task in the corresponding PE type total order. Then, any necessary actions are taken to maintain the priority ordering of descendant tasks.

[0167] An example pseudocode is shown below.

[0168] Algorithm Start

[0169] Input: Schedule (task set and PE total order set)

[0170] When a memory bottleneck exists in the scheduling process:

[0171] Step 1: Identify the earliest memory bottleneck in scheduling

[0172] Step 2: If there exists a bottleneck task (T1), and this task is not the last one in its corresponding PE total order (if there are multiple such tasks, then one is randomly selected), then

[0173]

[0174] otherwise

[0175] Correct the schedule and terminate execution.

[0176] Step 3: Rearrange T1 and its descendant tasks, as well as T2 and its descendant tasks, to obtain a new efficient schedule.

[0177] Step 4: Adjust the scheduling for tasks that started before the end of T2.

[0178] Step 5: Make the (remaining) schedule include only tasks that begin after T2 ends.

[0179] End the loop;

[0180] Adjust the scheduling;

[0181] Algorithm ends

[0182] Considering the examples in Figures 6(a) and 7(a), assuming a memory limit of 50, tasks 9 and 5 constitute the only bottleneck. Since there are no other tasks of the same type (hexagonal) after task 9, task 5 is swapped with the next task of the same type (elliptical), i.e., task 11. The tasks are rearranged as follows: due to priority constraints, task 6 is directly after task 5, while tasks 11-12 can be shifted so that task 11 is now directly after task 10. For the next iteration, if only tasks 12 and 5, which start after task 11, are considered, there are no other bottlenecks. The algorithm terminates and does indeed return a memory-feasible schedule for this input.

[0183] Figure 10 A flowchart illustrating the steps of an exemplary computer implementation of a method 1000 for processing a computational graph executed on a device, according to some embodiments of the present invention, is shown. The device includes a plurality of processing elements, each of which has access to shared memory.

[0184] In step 1001, the method includes receiving information about a plurality of processing elements of the device, each processing element being associated with a corresponding task type among a plurality of task types. In step 1002, the method includes receiving information indicating the size of the shared memory of the device. In step 1003, the method includes receiving a computation graph having a plurality of nodes corresponding to executable tasks and arcs indicating priority constraints between tasks, each task corresponding to a corresponding task type among a plurality of task types and constrained to execution in a processing element associated with its corresponding task type. In step 1004, the method includes receiving a plurality of original task full sequences, each original task full sequence corresponding to a corresponding processing element, wherein the plurality of original task full sequences adhere to priority constraints. In step 1005, the method includes determining a plurality of new task full sequences such that the peak memory usage for executing the computation graph does not exceed the size of the shared memory, each new task full sequence corresponding to a corresponding processing element, wherein the plurality of new task full sequences adhere to priority constraints.

[0185] This method is a static approach that makes minimal modifications to the task execution scheduling of multiple PEs to avoid memory shortages by adhering to device shared memory limits during parallel execution. It maintains the performance efficiency of a given schedule in an optimized manner while ensuring scheduling feasibility in terms of memory requirements. This method does not introduce any overhead as new synchronization events between PEs because it only introduces dependencies within each of the multiple tasks executed within the same PE. There is no need to optimize task ordering from scratch.

[0186] Previous methods did not consider explicit processing type constraints for the task. Operation reordering was only applied to a finite number of subgraphs (i.e., branch pairs) of the full neural network computation graph and / or determined from scratch, which is not necessary when a computationally efficient schedule is known in advance and it is desired to apply it to a new set of resources.

[0187] This method identifies the minimum set of task dependencies that need to be removed / introduced between tasks of the same type in a graph. The graph size remains asymptotically constant. There is no additional overhead from new synchronization events, and no redundant dependencies.

[0188] The described technical solution can also be applied to related fields where task scheduling needs to be adapted to sets of heterogeneous resources with limited memory. This method can be applied to situations with arbitrary priorities (modeled by a DAG) and processing type constraints. For example, such problems may arise in scheduling applications in real-time system design or parallel computing.

[0189] This method can be used as part of an AI compilation optimization approach. It can be reused in software systems focused on DNN execution. Examples include TensorFlow, PyTorch, JAX, XLA, and OpenXLA. The method can be implemented as part of a software platform or cloud service for DNN execution.

[0190] The applicant hereby separately discloses each individual feature described herein, as well as any combination of two or more such features. It is common knowledge to those skilled in the art that such features or combinations can be implemented as a whole according to this specification, regardless of whether such features or combinations of features solve any problem disclosed herein; and without limiting the scope of the claims. The applicant indicates that aspects of the invention can be constituted by any such individual features or combinations of features. In view of the foregoing description, various modifications that can be made within the scope of the invention will be apparent to those skilled in the art.

Claims

1. An apparatus (100) for processing a computational graph executed on a device, characterized in that, The device (150) includes multiple processing elements (151, 152, 153, 154), each of which can access shared memory (155). The device includes one or more processors (101) for: Receive (1001) information about the plurality of processing elements (151, 152, 153, 154) of the device, each processing element being associated with a corresponding task type among a plurality of task types; Receive (1002) information indicating the size of the shared memory (155) of the device (150); Receive (1003) a computation graph having multiple nodes corresponding to executable tasks and arcs indicating priority constraints between the tasks, each task corresponding to a corresponding task type among the multiple task types and subject to execution in the processing element (151, 152, 153, 154) associated with its corresponding task type. Receive (1004) multiple raw task full sequences, each raw task full sequence corresponding to a corresponding processing element (151, 152, 153, 154), wherein the multiple raw task full sequences comply with the priority constraint; Determine (1005) a plurality of new task full sequences such that the peak memory usage for executing the computation graph does not exceed the size of the shared memory (155), each new task full sequence corresponding to a corresponding processing element, wherein the plurality of new task full sequences comply with the priority constraints.

2. The apparatus (100) according to claim 1, characterized in that, The one or more processors (101) are used for: The multiple new tasks are converted into multiple non-redundant task dependencies in a total sequence. Insert the multiple non-redundant task dependencies into the computation graph.

3. The apparatus (100) according to claim 2, characterized in that, The one or more processors (101) are used for: Compile the computation graph; The compiled computation graph is sent to the device for execution via the plurality of processing elements and the shared memory.

4. The apparatus (100) according to claim 2 or 3, characterized in that, The one or more processors (101) are configured to convert the plurality of new tasks into a plurality of non-redundant task dependencies in the following manner: for each processing element, when traversing the corresponding new task total order of the corresponding processing element, for each pair of consecutive tasks, if there is no directed path of computation graph arc from the previous task to the next task in the corresponding pair of consecutive tasks, then a task dependency is defined.

5. The apparatus (100) according to any one of claims 2 to 4, characterized in that, The one or more processors (101) are used to insert the plurality of non-redundant task dependencies into the computation graph by introducing a new arc into the computation graph for each dependency.

6. The apparatus (100) according to any one of the preceding claims, characterized in that, The one or more processors (101) are used to determine the order of the plurality of new tasks in the following manner: Estimate the computation time and memory consumption for each task; The memory requirements during the execution of the computation graph are estimated based on the total sequence of the multiple original tasks; The total order of the plurality of new tasks is determined based on the estimated computation time and memory consumption of each task and the estimated memory requirements during the execution of the computation graph.

7. The apparatus (100) according to any one of the preceding claims, characterized in that, For each processing element (151, 152, 153, 154), the corresponding original task full sequence corresponds to the execution order of the plurality of tasks associated with the corresponding processing element.

8. The apparatus (100) according to any one of the preceding claims, characterized in that, The total sequence of the multiple new tasks has the smallest deviation from the total sequence of the multiple original tasks in a certain metric.

9. The apparatus (100) according to claim 8, characterized in that, The metric is the difference or relative error between the estimated computational graph execution time corresponding to the total sequence of the multiple original tasks and the total sequence of the multiple new tasks.

10. The apparatus (100) according to claim 8 or 9, characterized in that, The indicator is Kendall. Rank distance.

11. The apparatus (100) according to any one of the preceding claims, characterized in that, The computational graph represents a neural network.

12. The apparatus (100) according to any one of the preceding claims, characterized in that, The computation graph is a directed acyclic graph.

13. The apparatus (100) according to any one of the preceding claims, characterized in that, The computational graph is a layered graph.

14. A computer-implemented method (1000) for processing a computational graph executed on a device (150), characterized in that, The device includes multiple processing elements (151, 152, 153, 154), each of which can access shared memory (155), and the method includes: Receive (1001) information about the plurality of processing elements (151, 152, 153, 154) of the device (150), each processing element being associated with a corresponding task type among a plurality of task types; Receive (1002) information indicating the size of the shared memory (155) of the device (150); Receive (1003) a computation graph having multiple nodes corresponding to executable tasks and arcs indicating priority constraints between the tasks, each task corresponding to a corresponding task type among the multiple task types and restricted to execution in the processing element associated with its corresponding task type; Receive (1004) multiple raw task full sequences, each raw task full sequence corresponding to a corresponding processing element, wherein the multiple raw task full sequences comply with the priority constraint; Determine (1005) a plurality of new task full sequences such that the peak memory usage for executing the computation graph does not exceed the size of the shared memory, each new task full sequence corresponding to a corresponding processing element, wherein the plurality of new task full sequences comply with the priority constraints.

15. A computer program, characterized in that, Includes instructions that, when executed by a computer (100), cause the computer to perform the method (1000) according to claim 14.

Images