Task distribution method and device

Abstract

The invention provides a task distribution method and device, and the method comprises the steps: determining the number of writing subtasks based on the bottom data arrangement of an output tensor and the shape of the output tensor; determining the number of read sub-tasks based on the bottom data arrangement of the output tensor and the shape of the input tensor; after the read tasks are distributed based on the number of the read sub-tasks, write tasks are distributed based on the ratio of the number of the write sub-tasks to the number of the read sub-tasks. According to the method, the read-write sub-tasks can be reasonably divided according to the characteristics of underlying data arrangement, so that the adjacent sub-tasks access a continuous memory area as far as possible, in addition, a double-layer task distribution strategy is adopted, read task distribution is firstly performed based on the number of read sub-tasks, the priority of data reading is ensured, and the data reading efficiency is improved. And then writing task distribution is performed based on the number of the writing sub-tasks and the ratio of the number of the reading sub-tasks, so that the collaboration of the reading tasks and the writing tasks is ensured, and optimal configuration of computing resources is realized.

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence / chip technology, and in particular to a task distribution method and device. Background Art

[0002] In deep learning computations, data layout is crucial to operator performance. Different layouts influence data storage and access patterns, which in turn impacts memory access efficiency and computational performance. The goal of operator task distribution is to distribute computational tasks to different processing units (e.g., threads) as efficiently as possible while ensuring correctness, thereby fully utilizing hardware resources.

[0003] The current task distribution scheme first performs a shape expansion (expand) operation on the input tensor (src) and output tensor (dst) to expand src to the size of dst. The expanded src and dst are then distributed to multiple threads, each assigned a unique index. Based on this index, the offsets of the corresponding data in src and dst after the expand operation are independently calculated, and then the data is copied. However, this task distribution strategy ignores the complex layout changes that may be introduced by the expand operation. When operators involve non-contiguous or interleaved data layouts, the input tensor offsets (src_offset) and output tensor offsets (dst_offset) calculated by adjacent threads may show significant differences, resulting in the inability to merge memory read and write operations of adjacent threads, thereby preventing the effective use of continuous video memory access, significantly reducing memory bandwidth utilization, and becoming a performance bottleneck. Summary of the Invention

[0004] The present invention provides a task distribution method and device to solve the defects existing in the related art.

[0005] The present invention provides a task distribution method, comprising the following steps: Determining the number of write subtasks based on the underlying data layout of the output tensor and the shape of the output tensor; Determine the number of read subtasks based on the underlying data layout of the output tensor and the shape of the input tensor; After the read tasks are distributed based on the number of read subtasks, the write tasks are distributed based on the ratio between the number of write subtasks and the number of read subtasks.

[0006] According to a task distribution method provided by the present invention, the shape of the output tensor includes the dimensions of each dimension of the output tensor; The determining of the number of writing subtasks based on the underlying data arrangement of the output tensor and the shape of the output tensor includes: Determining a first dimension of the output tensor to be aligned and an alignment multiple of the first dimension to be aligned based on an underlying data arrangement of the output tensor; The number of the writing subtasks is determined based on the size of the first dimension to be aligned, the alignment multiple of the first dimension to be aligned, and the sizes of the remaining dimensions in the output tensor except the first dimension to be aligned.

[0007] According to a task distribution method provided by the present invention, determining the number of the writing subtasks based on the size of the first dimension to be aligned, the alignment multiple of the first dimension to be aligned, and the sizes of the remaining dimensions in the output tensor except the first dimension to be aligned includes: Based on the size of the first dimension to be aligned and the alignment multiple of the first dimension to be aligned, aligning on the first dimension to be aligned to determine a first alignment dimension size, where the first alignment dimension size refers to the dimension size after the first dimension to be aligned is aligned; The number of the writing subtasks is determined based on the first alignment dimension size and the remaining dimension sizes of the output tensor except the first to-be-aligned dimension.

[0008] According to a task distribution method provided by the present invention, the shape of the input tensor includes the dimensions of each dimension of the input tensor; Determining the number of read subtasks based on the underlying data arrangement of the output tensor and the shape of the input tensor includes: Determining a second dimension of the input tensor to be aligned and an alignment multiple of the second dimension based on an underlying data arrangement of the output tensor; The number of the read subtasks is determined based on the size of the second dimension to be aligned, the alignment multiple of the second dimension to be aligned, and the sizes of the remaining dimensions of the input tensor except the second dimension to be aligned.

[0009] According to a task distribution method provided by the present invention, determining the number of the read subtasks based on the size of the second dimension to be aligned, the alignment multiple of the second dimension to be aligned, and the remaining dimensions of the input tensor except the second dimension to be aligned includes: Based on the size of the second dimension to be aligned and the alignment multiple of the second dimension to be aligned, aligning on the second dimension to be aligned to determine a second alignment dimension size, where the second alignment dimension size refers to the dimension size after the second dimension to be aligned is aligned; The number of the reading subtasks is determined based on the size of the second alignment dimension and the sizes of the remaining dimensions in the output tensor except the second dimension to be aligned.

[0010] According to a task distribution method provided by the present invention, the read task reads data based on a data reading mode, and the data reading mode is determined based on the underlying data arrangement of the input tensor and the underlying data arrangement of the output tensor.

[0011] The present invention also provides a task distribution device, comprising the following modules: A first determining unit is configured to determine the number of writing subtasks based on an underlying data arrangement of an output tensor and a shape of the output tensor; A second determining unit is configured to determine the number of read subtasks based on the underlying data arrangement of the output tensor and the shape of the input tensor; The task distribution unit is configured to distribute the read tasks based on the number of the read subtasks, and then distribute the write tasks based on the ratio between the number of the write subtasks and the number of the read subtasks.

[0012] The present invention also provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the processor executes the program, any of the above-described task distribution methods is implemented.

[0013] The present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which implements any of the above-mentioned task distribution methods when executed by a processor.

[0014] The present invention also provides a computer program product, comprising a computer program, wherein when the computer program is executed by a processor, the computer program implements any of the above-mentioned task distribution methods.

[0015] The task distribution method and device provided by the present invention fully consider the underlying data layout of the output tensor, thereby being able to rationally divide read and write subtasks according to the characteristics of the underlying data layout, allowing adjacent subtasks to access continuous memory areas as much as possible, thereby greatly improving memory bandwidth utilization and reducing performance losses caused by non-contiguous memory access. At the same time, a two-tier task distribution strategy is adopted. Read tasks are first distributed based on the number of read subtasks to ensure the priority of data reading and maximize the reuse rate of data reading. Then, write tasks are distributed based on the ratio between the number of write subtasks and the number of read subtasks to ensure the synergy of read and write tasks, achieving optimal allocation of computing resources and avoiding resource waste or bottlenecks. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] In order to more clearly illustrate the technical solutions in the present invention or related technologies, the following is a brief introduction to the drawings required for use in the embodiments or related technical descriptions. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0017] Figure 1 It is a flowchart of the task distribution method provided by the present invention.

[0018] Figure 2 This is a schematic diagram of the underlying data arrangement provided by the present invention.

[0019] Figure 3 It is a structural diagram of the task distribution device provided by the present invention.

[0020] Figure 4 It is a structural schematic diagram of the electronic device provided by the present invention. DETAILED DESCRIPTION

[0021] To make the objectives, technical solutions, and advantages of the present invention more clear, the technical solutions of the present invention will be clearly and completely described below in conjunction with the accompanying drawings. Obviously, the embodiments described are only some of the embodiments of the present invention, not all of them. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts shall fall within the scope of protection of the present invention.

[0022] In order to facilitate understanding of the technical solution of the present invention, the relevant technical terms of the present invention are explained as follows: Layout: The underlying data arrangement is designed to accelerate the data read and write bandwidth of heterogeneous operators, thereby achieving higher performance. The underlying data arrangement can be row-major (also known as C-order), column-major (also known as Fortran-order), interleaved, tiled, or a custom complex layout.

[0023] Expand operation: The expand operator broadcasts the tensor to the specified size. Only the dimensions with an original size of 1 will be expanded. The other dimensions must match the size of the original tensor.

[0024] For example, assuming the shape of the expanded tensor is expand_size, for dimensions whose size in the input tensor is 1, the size of the corresponding dimension of expand_size can be set to a value greater than or equal to 1; for dimensions whose size in the input tensor is not 1, the size of the corresponding dimension of expand_size is equal to the size of the corresponding dimension in the input tensor.

[0025] The current task distribution scheme first performs a shape expansion (expand) operation on the input tensor (src) and output tensor (dst) to expand src to the size of dst. The expanded src and dst are then distributed to multiple threads, each assigned a unique index. Based on this index, the offsets of the corresponding data in src and dst after the expand operation are independently calculated, and then the data is copied. However, this task distribution strategy ignores the complex layout changes that may be introduced by the expand operation. When operators involve non-contiguous or interleaved data layouts, the input tensor offsets (src_offset) and output tensor offsets (dst_offset) calculated by adjacent threads may show significant differences, resulting in the inability to merge memory read and write operations of adjacent threads, thereby preventing the effective use of continuous video memory access, significantly reducing memory bandwidth utilization, and becoming a performance bottleneck.

[0026] For example, consider an expand operation that expands a tensor of shape (1,4,4) to (8,4,4). Suppose the original tensor uses a row-major layout, while the expanded tensor uses an interleaved layout due to some optimization strategy. In this case, even if the indices of adjacent threads differ by only 1, their memory addresses in the expanded tensor may be far apart. Thread 0 may need to access data at location (0,0,0), while thread 1 may need to access data at location (0,0,1). However, due to the interleaved layout, these two locations may not be physically adjacent in memory. This will cause each thread to independently read or write data from non-contiguous memory locations, making it impossible to merge memory access requests from adjacent threads. Ultimately, an operation that could have been completed with a single continuous read requires multiple non-contiguous reads, significantly reducing memory bandwidth utilization and thus affecting overall performance.

[0027] In addition, there are also customized optimization operators for each layout combination, and then task distribution for the customized optimized operators. However, customized optimization operators will cause the number of operators to grow exponentially with the number of layout combinations, which is extremely costly to implement and difficult to achieve generalization for layouts.

[0028] To this end, the present invention provides a task distribution method. Figure 1 It is a flowchart of the task distribution method provided by the present invention, such as Figure 1 As shown, the method includes step 110 , step 120 and step 130 .

[0029] Step 110: Determine the number of writing subtasks based on the underlying data layout of the output tensor and the shape of the output tensor.

[0030] Here, the output tensor can be thought of as the output of an operator, that is, the storage medium for the operator's calculation results. For example, for an expand operator, the output tensor is the result tensor after the expand operation. The underlying data layout of the output tensor refers to how the output tensor is stored in physical memory. This underlying data layout is used to represent the physical address relationship between each element in the output tensor. For example, if the underlying data layout of the output tensor is row-major layout, it means that in memory, elements in the same row of the output tensor are stored consecutively, while elements in different rows are stored sequentially according to row order.

[0031] The shape of the output tensor refers to the size information of the output tensor in each dimension, which includes the number of dimensions (rank) of the output tensor and the size of each dimension. For example, if the shape of the output tensor is (8,4,4), it means that the tensor is a three-dimensional tensor (rank=3), the size of the first dimension is 8, the size of the second dimension is 4, and the size of the third dimension is also 4.

[0032] As an optional embodiment, a generalized representation of the output tensor can be obtained from the output tensor's metadata or descriptor. This generalized representation can be understood as an abstract description of the output tensor's data layout, including information such as the output tensor's underlying data arrangement and shape. For example, for an output tensor with a shape of (8, 4, 4), if a row-major layout is used, its generalized representation may include (8, 4, 4) as the shape, and the underlying data arrangement is row-major.

[0033] A write subtask is an independent operation that writes a portion of the output tensor's data to memory, typically performed by a thread or a compute unit. The number of write subtasks is the total number of write subtasks required to complete the write operation of the entire output tensor.

[0034] Considering that traditional task distribution methods ignore the complex layout changes that may be introduced by the expand operation, resulting in discontinuous memory access between adjacent threads and inefficient use of video memory bandwidth, the embodiments of the present invention consider the layout of the output tensor when determining the number of write subtasks. This allows for a reasonable division of write subtasks based on the output tensor layout, ensuring that adjacent write subtasks access contiguous memory areas as much as possible, thereby improving memory access efficiency.

[0035] Furthermore, since each write subtask is responsible for writing a portion of the output tensor's data, and these write subtasks can be executed in parallel, the shape of the output tensor determines the maximum number of tasks that can be executed in parallel. For example, for an output tensor with a shape of (8, 4, 4), if the layout is not considered, each element can theoretically be processed as an independent write subtask, resulting in a maximum number of tasks that can be executed in parallel of 8 × 4 × 4 = 128.

[0036] To this end, an embodiment of the present invention combines the layout of the output tensor and the shape of the output tensor to determine the number of write subtasks, so that a continuous memory area can be divided into a write subtask according to the layout, reducing the synchronization overhead between threads and improving the utilization of memory bandwidth; at the same time, the number of write subtasks can also be controlled according to the shape of the output tensor to avoid excessive thread competition and improve computing efficiency.

[0037] Step 120: Determine the number of read subtasks based on the underlying data layout of the output tensor and the shape of the input tensor.

[0038] Specifically, the input tensor can be thought of as the input to an operator, the source of data required for the operator's computation. For example, for an expand operator, the input tensor is the original tensor to be expanded. The shape of an input tensor refers to the size of the input tensor in each dimension, including the number of dimensions (rank) and the size of each dimension.

[0039] In addition, a read subtask refers to the complete operation of reading the input tensor data from the memory for calculation, and the number of read subtasks refers to the total number of read subtasks that need to be divided in order to complete the reading operation of the entire input tensor.

[0040] In order to efficiently complete the filling of the output tensor, ensure that the operation of writing the output tensor is as continuous and efficient as possible. To achieve this, it is necessary to arrange the reading of data according to the layout of the output tensor, so as to ensure that the read data can be written to the output tensor in the most suitable way. If the read subtask is determined based on the layout of the input tensor, it may cause the read data to need to be rearranged when writing to the output tensor, which will reduce efficiency. In addition, step 110 has determined the number of write subtasks based on the layout of the output tensor. In order to ensure the coordination of read and write operations, step 120 needs to determine the read subtask based on the layout of the output tensor, so that each read subtask can provide the required data for the corresponding write subtask, which can avoid the mismatch between reading and writing and reduce the overhead of data transmission and conversion.

[0041] Furthermore, the shape of the input tensor defines the valid boundaries of the data. The read subtask must operate within this boundary, otherwise it will cause out-of-bounds access and generate an error. The shape of the input tensor can ensure that each read subtask only reads valid data. In addition, the read subtask needs to know where to start reading the data in the input tensor. The shape of the input tensor is an important basis for calculating these index offsets. For example, if the input tensor is a three-dimensional array with a shape of (D1, D2, D3), then reading the (i, j, k)th element requires calculating the correct memory offset based on this shape.

[0042] Based on this, the embodiment of the present invention combines the underlying data arrangement of the output tensor and the shape of the input tensor to determine the number of read subtasks, so as to maximize the memory access efficiency while satisfying the data dependency, and achieve efficient collaboration with the write subtasks, ultimately improving the overall performance of the operator.

[0043] Step 130 : After distributing the read tasks based on the number of read subtasks, distribute the write tasks based on the ratio between the number of write subtasks and the number of read subtasks.

[0044] Specifically, the write subtask is an independent operation that writes part of the output tensor data into the memory, and the read subtask is an independent operation that reads part of the input tensor data from the memory for calculation. In the operator calculation process, the output tensor is obtained after the operation is performed based on the input tensor, that is, the operation of reading the input tensor is performed first, and then the operation of writing the output tensor is performed. Since the execution of the write operation depends on the result of the read operation, the read subtask is the prerequisite for the write subtask. On this basis, the embodiment of the present invention first distributes the read task and then distributes the write task, so as to ensure that the required data is ready before the write operation is performed, avoid data dependency problems, ensure the correctness of the calculation, and ultimately improve the overall efficiency of the operator.

[0045] Furthermore, by prioritizing all read subtasks across different resources, we ensure that each resource is responsible for reading a portion of the data as independently as possible. Subsequent write subtasks can loop over one or more resources, repeatedly utilizing data previously read by the read subtask. This means that data only needs to be read once to be used by multiple write subtasks, avoiding duplicate reads and improving memory bandwidth utilization.

[0046] After determining the number of read subtasks, the read tasks can be distributed according to available resources (such as warp resources, i.e., parallel computing units on the GPU). For example, if the number of read subtasks is less than or equal to the number of available resources, each read subtask can be assigned an independent available resource; if the number of read subtasks is greater than the number of available resources, multiple read subtasks can be assigned to the same available resource.

[0047] The ratio between the number of write subtasks and the number of read subtasks directly reflects the number of write subtasks each read subtask needs to provide data for. If the ratio is greater than 1, it means that a read subtask needs to service multiple write subtasks. Based on this ratio, allocating write subtasks can minimize the read and write load on each resource (such as a thread). If a thread is responsible for too many write tasks, it may become a bottleneck; conversely, if a thread is responsible for too few write tasks, resources may be wasted.

[0048] For example, if the number of read subtasks is 4 and the number of write subtasks is 8, the read-write ratio is 2. If the total number of resources is 4, each resource is responsible for executing one read subtask and two write subtasks at the same time. While ensuring that resources are not wasted and the load is balanced, each read data is reused twice to achieve the maximum data reuse rate.

[0049] If the total number of resources is 8, the resources are divided into 4 × 2 parts. The resources are first allocated to the read subtask, and then the remaining resource ratio 2 is divided into the read and write ratio. In this way, each resource is responsible for one read subtask and one write subtask. While ensuring that resources are not wasted and the load is balanced, each read data is reused once.

[0050] More complicatedly, if the total number of resources is 6, without changing the number of read and write subtasks, the resources can only be divided into 4×1 parts, which is equivalent to the case where the total number of resources is 4. Here, due to the mismatch between the total number of resources and the number of tasks, a small amount of resources are wasted. Here, the resource waste problem can be solved by adjusting the number of read and write subtasks while ensuring minimum continuous access. For example, the number of read subtasks can be divided into 3 and the number of write subtasks can be divided into 6. The read-write ratio can always be 2, and the resources can be divided into 3×2 parts. This ensures that each read data is reused twice without wasting resources and load balancing, achieving the maximum data reuse rate.

[0051] The task distribution method provided by an embodiment of the present invention fully considers the underlying data layout of the output tensor, thereby rationally dividing read and write subtasks according to the characteristics of the layout. This allows adjacent subtasks to access contiguous memory areas as much as possible, significantly improving memory bandwidth utilization and reducing performance losses caused by discontinuous memory access. Furthermore, a two-tier task distribution strategy is employed: first, read tasks are distributed based on the number of read subtasks to ensure data read priority and maximize data read reuse, and then, write tasks are distributed based on the ratio between the number of write and read subtasks to ensure the synergy of read and write tasks, achieving optimal allocation of computing resources and avoiding resource waste or bottlenecks.

[0052] In an embodiment of the present invention, the above-mentioned operators can be regarded as components in a neural network model. The neural network model can be applied to fields such as text processing, image processing, and speech processing. The neural network model can, for example, be a convolutional neural network model (CNN), a recurrent neural network model (RNN), etc.

[0053] Taking the application of neural network models to text processing as an example, the aforementioned operators can be any parameters used, input, or generated in tasks such as text classification, sentiment analysis, and text generation, such as the semantic feature vector of the text. After executing the task distribution method based on the above embodiment, the input in each thread is specifically transformed to obtain a calculation result. The model then performs inference based on the calculation result, processes the input text, and obtains the corresponding output text.

[0054] Taking the application of a neural network model to image processing as an example, the target operator can be any parameter used, input, or generated in tasks such as image classification, object detection, and image segmentation, such as an image feature map. After task distribution is performed based on the method of the above embodiment, specific transformations are performed on the inputs in each thread to obtain computational results. The model then performs inference based on these computational results, processes the input image, and obtains corresponding output results (e.g., classification labels, bounding boxes, segmentation masks).

[0055] Taking the application of a neural network model to speech processing as an example, the target operator can be any parameter used, input, or generated in tasks such as speech recognition, speech synthesis, and speech enhancement, such as speech spectral features (Mel-spectrogram, MFCC). After task distribution based on the method of the above embodiment, the input speech signal in each thread is processed to obtain the corresponding output text or synthesized speech.

[0056] Based on the above embodiment, the shape of the output tensor includes the size of each dimension of the output tensor; Based on the underlying data layout of the output tensor and the shape of the output tensor, the number of write subtasks is determined, including: Determine a first dimension of the output tensor to be aligned and an alignment multiple of the first dimension based on the underlying data arrangement of the output tensor. The number of write subtasks is determined based on the size of the first dimension to be aligned, the alignment multiple of the first dimension to be aligned, and the sizes of the remaining dimensions in the output tensor except the first dimension to be aligned.

[0057] Specifically, the first dimension to be aligned refers to the dimension that needs to be aligned on the output tensor when determining the number of write subtasks. For example, if the shape of the output tensor is (n, h, w), the underlying data is arranged to combine the 4 data on the h dimension and the 32 data on the w dimension to form a continuous memory block of size 4×32=128, and then arrange the w, h, and n dimensions in sequence. The first alignment dimension includes the h dimension and the w dimension.

[0058] In addition, the alignment multiple of the first dimension to be aligned refers to the least common multiple to which the first dimension to be aligned needs to be aligned. This alignment multiple can be determined based on the underlying data layout. For example, if the underlying data layout combines 4 data in the h dimension and 32 data in the w dimension to form a contiguous memory block of size 4×32=128, then the h dimension is aligned to a multiple of 4, and the w dimension is aligned to a multiple of 32. In other words, the h dimension alignment multiple is 4, and the w dimension alignment multiple is 32.

[0059] The alignment multiple of the first dimension to be aligned determines the minimum granularity of the write subtask. For example, if the alignment multiple of the h dimension is 4, each write subtask must process at least 4 rows of data. The size of the first dimension to be aligned determines how many write subtasks can be divided along that dimension while still meeting the alignment requirements.

[0060] In addition to the first dimension to be aligned, the output tensor also includes the remaining dimensions, which do not require alignment. When calculating the number of write subtasks, in addition to considering the size of the first dimension to be aligned and the alignment multiple of the first dimension to be aligned, the sizes of the remaining dimensions are also used to determine whether to further split the write subtasks along these dimensions to improve parallelism and balance the load. In other words, the larger the remaining dimensions, the more write subtasks can be split.

[0061] For example, the output tensor shape is (n, h, w) = (8, 10, 60). The underlying data layout requires the h dimension to be aligned to a multiple of 4 and the w dimension to be aligned to a multiple of 32. Based on this, it can be determined that the h dimension needs to be aligned to 12 and the w dimension needs to be aligned to 64, which means that the output tensor shape after alignment is (8, 12, 64). In the h dimension, there are 12 / 4 = 3 write subtasks; in the w dimension, there are 64 / 32 = 2 write subtasks. The n dimension size is 8, which means that tasks can still be divided in the n dimension. For example, if each (h, w) block is divided equally into 8 parts along the n dimension, there will be a total of 3 × 2 × 8 = 48 write subtasks.

[0062] Based on any of the foregoing embodiments, determining the number of write subtasks based on the size of the first dimension to be aligned, the alignment multiple of the first dimension to be aligned, and the sizes of the remaining dimensions in the output tensor except the first dimension to be aligned includes: Based on the size of the first dimension to be aligned and the alignment multiple of the first dimension to be aligned, alignment is performed on the first dimension to be aligned to determine a first alignment dimension size, where the first alignment dimension size refers to the dimension size after the first dimension to be aligned is aligned; The number of write subtasks is determined based on the size of the first alignment dimension and the sizes of the remaining dimensions in the output tensor except the first to-be-aligned dimension.

[0063] Specifically, the size of the first alignment dimension refers to the size of the corresponding alignment dimension after the output tensor is dimensionally aligned. The size of the first alignment dimension is an integer multiple of the corresponding alignment multiple and is greater than or equal to the size of the first dimension to be aligned. When aligning on the first dimension to be aligned, you can choose a minimum alignment multiple greater than or equal to the size of the first dimension to be aligned, or you can choose a larger alignment multiple to meet specific hardware or performance requirements. For example, the output tensor shape is (n,h,w)=(8,10,60), and the underlying data layout requires the h dimension to be aligned to a multiple of 4 and the w dimension to be aligned to a multiple of 32. You can either align the h dimension to 12 and the w dimension to 64 (that is, choose a minimum alignment multiple greater than or equal to the size of the first dimension to be aligned), or you can align the h dimension to 16 and the w dimension to 96 (that is, choose a larger alignment multiple).

[0064] Since the size of the first alignment dimension is the size after the alignment along the first to-be-aligned dimension, the shape of the new output tensor can be determined by combining the size of the first alignment dimension and the sizes of the remaining dimensions. This shape is the shape of the output tensor after the alignment along the first to-be-aligned dimension. Based on the shape of the new output tensor, different task partitioning strategies can be selected. For example, each element can be treated as a separate write subtask, each row or column as a write subtask, or the entire tensor as a write subtask, which can then determine the number of write subtasks. For example, if the shape of the new output tensor is (n, h_align, w_aligned) = (8, 12, 64), the number of write subtasks can be selected based on different strategies. For example, if each element is a separate write subtask, the number of write subtasks is 8 × 12 × 64 = 6144; if each row is a write subtask, the number of write subtasks is 8 × 12 = 96; and if the entire tensor is a write subtask, the number of write subtasks is 1.

[0065] Taking an interleave layout as an example, the shape of the output tensor is (n, h, w), where the underlying data is arranged as a 4×32 block of continuous memory in the h and w dimensions, and then the w, h, and n dimensions are arranged in sequence. Figure 2 This is a schematic diagram of the underlying data arrangement provided by the present invention, such as Figure 2As shown, each small square contains 4 data in the h dimension and 32 data in the w dimension, that is, 4×32=128 data. After alignment in the h and w dimensions, the shape of the output tensor is (n,h_align4,w_align32). The number of write subtasks store_nums can be calculated based on the following formula: store_nums=n×num_h×num_w Among them, num_w = w_align32 / 32, num_h = h_align4 / unroll_h. Considering the flexibility of calculation and optimization of task scheduling, unroll_h can be any value among 1, 2, 3, and 4, which can be adjusted according to the task amount.

[0066] Based on any of the above embodiments, the shape of the input tensor includes the size of each dimension of the input tensor; Based on the underlying data layout of the output tensor and the shape of the input tensor, the number of reading subtasks is determined, including: Determine the second dimension of the input tensor to be aligned and the alignment multiple of the second dimension based on the underlying data arrangement of the output tensor. The number of read subtasks is determined based on the size of the second dimension to be aligned, the alignment multiple of the second dimension to be aligned, and the sizes of the remaining dimensions in the input tensor except the second dimension to be aligned.

[0067] Specifically, the second dimension to be aligned refers to the dimension that needs to be aligned on the input tensor when determining the number of read subtasks. In addition, the alignment multiple of the second dimension to be aligned refers to the least common multiple to which the second dimension to be aligned needs to be aligned, and the alignment multiple can be determined based on the underlying data arrangement. For example, the underlying data arrangement is to combine the 4 data on the h dimension and the 32 data on the w dimension to form a continuous memory block of size 4×32=128, then align the h dimension to a multiple of 4, and align the w dimension to a multiple of 32, that is, the alignment multiple of the h dimension is 4, and the alignment multiple of the w dimension is 32.

[0068] The alignment multiple of the second dimension determines the minimum granularity of the read subtask. For example, if the alignment multiple of the h dimension is 4, each read subtask must process at least 4 rows of data. The size of the second dimension determines how many read subtasks can be divided along that dimension while still meeting the alignment requirements.

[0069] In addition to the second dimension to be aligned, the input tensor also includes the remaining dimensions, which do not require alignment. When calculating the number of read subtasks, in addition to considering the size of the second dimension to be aligned and the alignment multiple of the second dimension to be aligned, the sizes of the remaining dimensions are also used to determine whether to further split the read subtasks along these dimensions to improve parallelism and balance the load. In other words, the larger the remaining dimensions, the more read subtasks can be split.

[0070] Based on any of the foregoing embodiments, determining the number of read subtasks based on the size of the second dimension to be aligned, the alignment multiple of the second dimension to be aligned, and the sizes of the remaining dimensions in the input tensor except the second dimension to be aligned includes: Based on the size of the second dimension to be aligned and the alignment multiple of the second dimension to be aligned, aligning on the second dimension to be aligned to determine the size of the second alignment dimension, where the size of the second alignment dimension refers to the size of the dimension after the second dimension to be aligned is aligned; The number of read subtasks is determined based on the size of the second alignment dimension and the sizes of the remaining dimensions in the output tensor except the second to-be-aligned dimension.

[0071] Specifically, the size of the second alignment dimension refers to the size of the corresponding alignment dimension after the input tensor is dimensionally aligned. The size of the second alignment dimension is an integer multiple of the corresponding alignment multiple and is greater than or equal to the size of the second dimension to be aligned. When aligning on the second dimension to be aligned, you can choose a minimum alignment multiple greater than or equal to the size of the second dimension to be aligned, or you can choose a larger alignment multiple to meet specific hardware or performance requirements. For example, the input tensor shape is (n,h,w)=(8,10,60), and the underlying data arrangement requires the h dimension to be aligned to a multiple of 4 and the w dimension to be aligned to a multiple of 32. You can either align the h dimension to 12 and the w dimension to 64 (that is, choose a minimum alignment multiple greater than or equal to the size of the second dimension to be aligned), or you can align the h dimension to 16 and the w dimension to 96 (that is, choose a larger alignment multiple).

[0072] Since the second alignment dimension size is the size after alignment along the second to-be-aligned dimension, the shape of the new input tensor can be determined by combining the size of the second alignment dimension and the sizes of the remaining dimensions. This shape is the shape of the input tensor after alignment along the second to-be-aligned dimension. Based on the shape of the new input tensor, different task partitioning strategies can be selected. For example, each element can be treated as a separate read subtask, each row or column as a read subtask, or the entire tensor as a read subtask, which can then determine the number of read subtasks. For example, if the shape of the new input tensor is (n, h_align, w_aligned) = (8, 12, 64), the number of read subtasks can be selected based on different strategies. For example, if each element is a separate read subtask, the number of read subtasks is 8 × 12 × 64 = 6144; if each row is a read subtask, the number of read subtasks is 8 × 12 = 96; and if the entire tensor is a read subtask, the number of read subtasks is 1.

[0073] Assuming the shape of the new input tensor is (n, h_align4, w_align32), the number of read subtasks load_nums can be calculated based on the following formula: load_nums=n×num_h×num_w Among them, num_w = w_align32 / 32, num_h = h_align4 / unroll_h. Considering the flexibility of calculation and optimization of task scheduling, unroll_h' can be any value among 1, 2, 3, and 4, which can be adjusted according to the task amount.

[0074] Based on any of the above embodiments, the read task reads data based on a data reading mode, and the data reading mode is determined based on the underlying data arrangement of the input tensor and the underlying data arrangement of the output tensor.

[0075] Specifically, the data reading mode refers to the method of reading data from the input tensor to meet the computational requirements of the output tensor. Different data reading modes are required for different combinations of the underlying data layout of the input and output tensors to ensure data reading efficiency and correctness, and to maximize memory bandwidth utilization. For example, if the input and output tensors have the same underlying data layout, the merge reading mode can be used; if the input tensor needs to be broadcast, the data broadcast mode can be used.

[0076] Among them, the data reading mode includes at least one of the flexible reading mode, the data broadcasting mode and the merged reading mode. The flexible reading mode means that each thread independently calculates the required memory address and loads the data from the memory separately. This mode has the highest flexibility and can handle various complex memory layouts, but the efficiency is relatively low. The data broadcasting mode means that one thread (usually thread 0) loads the data required by all threads, and then distributes the data to other threads through the thread broadcast mechanism within the warp. This mode is suitable for scenarios where the input tensor needs to be broadcast within the warp, which can reduce the number of memory accesses. The merged reading mode means that all threads within the warp work together to merge the data to be read into one or several large memory blocks, and then load them into the shared memory at one time. This mode is suitable for scenarios where the input tensors are stored continuously in the memory, which can maximize the memory bandwidth utilization.

[0077] Taking an interleave layout as an example, the shape of the output tensor is (n, h, w). The 32 data in the underlying w dimension are arranged continuously, so 32 threads can merge and write back data in the w dimension.

[0078] a) If the input tensor has a shape of (n, 1, w) and the layout of the input tensor is not continuous with 32 data points along the w dimension, each thread in the warp needs to read 32 different data points adjacent to h = 0 and w. Since the input tensor is not continuous along the w dimension, each thread can only calculate the corresponding address, that is, use flexible read mode.

[0079] b) If the shape of the input tensor is (n, h, 1), that is, broadcasting on the last dimension, thread 0 can read the data at position w=0 and then broadcast it to other threads in the warp, regardless of whether the layout of the input tensor is the same as the layout of the output tensor. This is the data broadcast mode.

[0080] c) If the shape of the input tensor is (n, 1, w), and the layout of the input tensor is the same as the layout of the output tensor, then each thread in the warp needs to read 32 different adjacent data with h = 0 and w. Since the input is also arranged continuously in the w dimension, the 32 adjacent data can be merged and read, that is, the merged read mode is adopted.

[0081] It can be seen that the embodiments of the present invention flexibly select data reading modes, thereby making full use of hardware characteristics (such as the thread broadcast mechanism within a warp) and selecting the optimal reading method based on the layout characteristics of the input tensors and output tensors to maximize memory bandwidth utilization and reduce unnecessary memory access times, ultimately improving the overall performance of the operator.

[0082] The task distribution device provided by the present invention is described below. The task distribution device described below and the task distribution method described above can be referenced to each other.

[0083] Based on any of the above embodiments, Figure 3 It is a structural diagram of the task distribution device provided by the present invention, such as Figure 3 As shown, the device includes: A first determining unit 310 is configured to determine the number of writing subtasks based on the underlying data arrangement of the output tensor and the shape of the output tensor; A second determining unit 320 is configured to determine the number of read subtasks based on the underlying data arrangement of the output tensor and the shape of the input tensor; The task distribution unit 330 is configured to distribute the read tasks based on the number of read subtasks, and then distribute the write tasks based on the ratio between the number of write subtasks and the number of read subtasks.

[0084] Based on any of the above embodiments, the shape of the output tensor includes the size of each dimension of the output tensor; Based on the underlying data layout of the output tensor and the shape of the output tensor, the number of write subtasks is determined, including: Determine a first dimension of the output tensor to be aligned and an alignment multiple of the first dimension based on the underlying data arrangement of the output tensor. The number of write subtasks is determined based on the size of the first dimension to be aligned, the alignment multiple of the first dimension to be aligned, and the sizes of the remaining dimensions in the output tensor except the first dimension to be aligned.

[0085] Based on any of the foregoing embodiments, determining the number of write subtasks based on the size of the first dimension to be aligned, the alignment multiple of the first dimension to be aligned, and the sizes of the remaining dimensions in the output tensor except the first dimension to be aligned includes: Based on the size of the first dimension to be aligned and the alignment multiple of the first dimension to be aligned, alignment is performed on the first dimension to be aligned to determine a first alignment dimension size, where the first alignment dimension size refers to the dimension size after the first dimension to be aligned is aligned; The number of write subtasks is determined based on the size of the first alignment dimension and the sizes of the remaining dimensions in the output tensor except the first to-be-aligned dimension.

[0086] Based on any of the above embodiments, the shape of the input tensor includes the size of each dimension of the input tensor; Based on the underlying data layout of the output tensor and the shape of the input tensor, the number of reading subtasks is determined, including: Determine the second dimension of the input tensor to be aligned and the alignment multiple of the second dimension based on the underlying data arrangement of the output tensor. The number of read subtasks is determined based on the size of the second dimension to be aligned, the alignment multiple of the second dimension to be aligned, and the sizes of the remaining dimensions in the input tensor except the second dimension to be aligned.

[0087] Based on any of the foregoing embodiments, determining the number of read subtasks based on the size of the second dimension to be aligned, the alignment multiple of the second dimension to be aligned, and the sizes of the remaining dimensions in the input tensor except the second dimension to be aligned includes: Based on the size of the second dimension to be aligned and the alignment multiple of the second dimension to be aligned, aligning on the second dimension to be aligned to determine the size of the second alignment dimension, where the size of the second alignment dimension refers to the size of the dimension after the second dimension to be aligned is aligned; The number of read subtasks is determined based on the size of the second alignment dimension and the sizes of the remaining dimensions in the output tensor except the second to-be-aligned dimension.

[0088] Based on any of the above embodiments, the read task reads data based on a data reading mode, and the data reading mode is determined based on the underlying data arrangement of the input tensor and the underlying data arrangement of the output tensor.

[0089] Figure 4 Schematic diagram of the structure of the electronic device provided by the present invention, such as Figure 4 As shown, the electronic device may include: a processor 410, a communications interface 420, a memory 430, and a communications bus 440, wherein the processor 410, the communications interface 420, and the memory 430 communicate with each other via the communications bus 440. The processor 410 may call logic instructions in the memory 430 to execute a task distribution method, which includes: determining the number of write subtasks based on the underlying data arrangement of the output tensor and the shape of the output tensor; determining the number of read subtasks based on the underlying data arrangement of the output tensor and the shape of the input tensor; after distributing read tasks based on the number of read subtasks, distributing write tasks based on the ratio between the number of write subtasks and the number of read subtasks.

[0090] Furthermore, the logic instructions in the aforementioned memory 430 can be implemented as software functional units and, when sold or used as independent products, can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the portion that contributes to the relevant art, or a portion of the technical solution, can be embodied in the form of a software product. This computer software product, stored in a storage medium, includes instructions for enabling a computer device (which can be a personal computer, server, or network device, etc.) to perform all or part of the steps of the methods described in various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as a USB flash drive, a mobile hard drive, a read-only memory (ROM), a random access memory (RAM), a magnetic disk, or an optical disk.

[0091] On the other hand, the present invention also provides a computer program product, which includes a computer program, which can be stored on a non-transitory computer-readable storage medium. When the computer program is executed by a processor, the computer can execute the task distribution method provided by the above methods, which includes: determining the number of write subtasks based on the underlying data arrangement of the output tensor and the shape of the output tensor; determining the number of read subtasks based on the underlying data arrangement of the output tensor and the shape of the input tensor; after distributing read tasks based on the number of read subtasks, distributing write tasks based on the ratio between the number of write subtasks and the number of read subtasks.

[0092] On the other hand, the present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, is implemented to execute the task distribution method provided by the above-mentioned methods, the method comprising: determining the number of write subtasks based on the underlying data arrangement of the output tensor and the shape of the output tensor; determining the number of read subtasks based on the underlying data arrangement of the output tensor and the shape of the input tensor; after distributing read tasks based on the number of read subtasks, distributing write tasks based on the ratio between the number of write subtasks and the number of read subtasks.

[0093] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units, i.e., they may be located in one location or distributed across multiple network units. Some or all of the modules may be selected based on actual needs to achieve the objectives of the present embodiment. Persons of ordinary skill in the art will be able to understand and implement the present invention without inventive effort.

[0094] Through the description of the above embodiments, those skilled in the art will clearly understand that each embodiment can be implemented using software plus a necessary general-purpose hardware platform, or of course, hardware. Based on this understanding, the essence of the above technical solution, or the portion that contributes to the relevant technology, can be embodied in the form of a software product. This computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, a magnetic disk, or an optical disk, and includes a number of instructions for causing a computer device (such as a personal computer, server, or network device) to execute the methods described in each embodiment or certain portions of the embodiments.

[0095] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the various embodiments of the present invention.

Claims

1. A task distribution method, characterized in that: include: Determining the number of write subtasks based on the underlying data layout of the output tensor and the shape of the output tensor; Determine the number of read subtasks based on the underlying data layout of the output tensor and the shape of the input tensor; After the read tasks are distributed based on the number of read subtasks, the write tasks are distributed based on the ratio between the number of write subtasks and the number of read subtasks.

2. The task distribution method according to claim 1, characterized in that: The shape of the output tensor includes the dimensions of each dimension of the output tensor; The determining of the number of writing subtasks based on the underlying data arrangement of the output tensor and the shape of the output tensor includes: Determining a first dimension of the output tensor to be aligned and an alignment multiple of the first dimension to be aligned based on an underlying data arrangement of the output tensor; The number of the writing subtasks is determined based on the size of the first dimension to be aligned, the alignment multiple of the first dimension to be aligned, and the sizes of the remaining dimensions in the output tensor except the first dimension to be aligned.

3. The task distribution method according to claim 2, characterized in that: The determining the number of the writing subtasks based on the size of the first dimension to be aligned, the alignment multiple of the first dimension to be aligned, and the sizes of the remaining dimensions in the output tensor except the first dimension to be aligned includes: Based on the size of the first dimension to be aligned and the alignment multiple of the first dimension to be aligned, aligning on the first dimension to be aligned to determine a first alignment dimension size, where the first alignment dimension size refers to the dimension size after the first dimension to be aligned is aligned; The number of the writing subtasks is determined based on the first alignment dimension size and the remaining dimension sizes of the output tensor except the first to-be-aligned dimension.

4. The task distribution method according to any one of claims 1 to 3, characterized in that: The shape of the input tensor includes the dimensions of each dimension of the input tensor; Determining the number of read subtasks based on the underlying data arrangement of the output tensor and the shape of the input tensor includes: Determining a second dimension of the input tensor to be aligned and an alignment multiple of the second dimension based on an underlying data arrangement of the output tensor; The number of the read subtasks is determined based on the size of the second dimension to be aligned, the alignment multiple of the second dimension to be aligned, and the sizes of the remaining dimensions of the input tensor except the second dimension to be aligned.

5. The task distribution method according to claim 4, characterized in that: The determining the number of the read subtasks based on the size of the second dimension to be aligned, the alignment multiple of the second dimension to be aligned, and the remaining dimensions of the input tensor except the second dimension to be aligned includes: Based on the size of the second dimension to be aligned and the alignment multiple of the second dimension to be aligned, aligning on the second dimension to be aligned to determine a second alignment dimension size, where the second alignment dimension size refers to the dimension size after the second dimension to be aligned is aligned; The number of the reading subtasks is determined based on the size of the second alignment dimension and the sizes of the remaining dimensions in the output tensor except the second dimension to be aligned.

6. The task distribution method according to any one of claims 1 to 3, characterized in that: The read task reads data based on a data reading mode, where the data reading mode is determined based on an underlying data arrangement of the input tensor and an underlying data arrangement of the output tensor.

7. A task distribution device, characterized in that: include: A first determining unit is configured to determine the number of writing subtasks based on an underlying data arrangement of an output tensor and a shape of the output tensor; A second determining unit is configured to determine the number of read subtasks based on the underlying data arrangement of the output tensor and the shape of the input tensor; The task distribution unit is configured to distribute the read tasks based on the number of the read subtasks, and then distribute the write tasks based on the ratio between the number of the write subtasks and the number of the read subtasks.

8. An electronic device comprising a memory, a processor, and a computer program stored in the memory and running on the processor, characterized in that: When the processor executes the computer program, the task distribution method according to any one of claims 1 to 6 is implemented.

9. A non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the task distribution method according to any one of claims 1 to 6 is implemented.

10. A computer program product comprising a computer program, characterized in that When the computer program is executed by a processor, the task distribution method according to any one of claims 1 to 6 is implemented.

Images