Execution Process

Single-operator Execution Process

The execution process of a single-operator in the ATB consists of the following steps: kernel graph building, required input for kernel running, device memory computation, tiling data computing and moving, and computation task delivery.

1. Kernel graph building

   A kernel is a basic code unit running on a device, and is similar to a function in C language. On the device, various computing tasks are executed by using the kernel as a unit. Kernel graph building is essentially building of a kernel task queue, or building of a device task flow.

   The operations provided by the ATB are kernel combinations with complex functions. The relationships between kernels are organized in the form of a graph. The same operation combines different kernels when using different features or inputs. Therefore, the kernel graph or device task flow corresponding to a single operation can be determined only during running.

2. Necessary input for kernel running

   To ensure that a kernel task flow can be properly executed, you need to prepare the corresponding input for each kernel on the device.

   For a single kernel, there are three necessary input types: input and output tensors, tiling data, and scratch memory. The three types of inputs are provided to the kernel in the form of device-side addresses.

   • The input and output tensors are provided by users. The ATB prepares the tiling data and scratch memory required by the kernel.
   • Tiling data is a kernel tiling parameter, which is used to determine the actual tiling policy during kernel computation. It is usually stored in a structure and calculated based on the input parameters and tensor shape.
   • The scratch memory is the space used by the kernel to store temporary data.
3. Device memory calculation

   For a kernel graph, after the input and output tensors of the entire graph are provided, the intermediate tensors of the graph need to be allocated. Therefore, in addition to the tiling data and scratch memory of each kernel, the ATB also prepares the intermediate tensor of the kernel graph.

   A single kernel can be used as a kernel graph with the intermediate tensor size of 0. The ATB calculates the device memory required by a single-operator based on the kernel graph.

   The ATB calculates the intermediate tensor size required by the kernel graph and the scratch memory size of each kernel, and then returns the values as the WorkSpaceSize to the user. The tiling data is allocated and managed by the ATB in a unified manner after the size is calculated.

4. Tiling data calculation and moving

   The tiling data is usually computed on the host. After the computation is complete on the host, the ATB copies the tiling data to the device as the input of the kernel.

5. Computing task delivery

   After the task execution queue is determined and the input is ready, the last step is to deliver the task to the device.

   In this step, the ATB uses the prepared input as the input parameter to the task delivery API of the device based on the constructed kernel queue, and waits for the task to be executed on the device.

Tiling Data Moving Policies

The ATB supports three tiling data moving policies: overall, multi-stream, and kernel-dependent. The current setup defaults to moving with the kernel.

• In overall tiling moving mode, the computed tiling data is stored in a contiguous host memory and moved to the device at a time after the tiling data computation of all kernels is complete.
• The multi-stream moving mode is an improvement on the overall moving mode. The core idea is to reduce the time consumed by tiling data moving in parallel stream mode.

  In this moving policy, the ATB prepares two streams, a ring device buffer, and a series of synchronization semaphores.

  • One stream is used for kernel execution, and the other stream is used for copying tiling data. The kernel execution depends on the completion of tiling data copy. In this case, you need to use the synchronization semaphore to ensure that the kernel execution action on the other stream is performed after the tiling data copy of the current kernel is complete.
  • The ring device memory buffer is used to store the tiling data that is copied to the device in advance when the tiling data copy is faster than the kernel execution, so that the tiling data does not conflict with each other.
    However, when the tiling data copy is too fast, the device buffer may be full. In this case, you need to increase the number of memory blocks in the device buffer.

    

    This function is supported only for Atlas inference products.

  Figure 1 Multi-Stream moving policy
  
• Tiling data moving with the kernel is the default tiling data moving policy.

  This policy optimizes the performance of the overall tiling moving mode, but the optimization mode is different from the multi-stream moving mode.

  The core idea is as follows:

  • The tiling data is not moved to the device at a time after all tiling data is calculated. Instead, the tiling data of a kernel is moved to the device each time the tiling data of the kernel is calculated.
  • When the kernel task is delivered to the device, the tiling data moving task is started at the same time. In this way, devices on the host and device can run in parallel. When the host prepares the tiling data of the next kernel, the device executes the current kernel task at the same time, significantly improving the tiling data moving efficiency.

  As shown in Figure 2, compared with the multi-stream moving mode, when the tiling data moving is too fast, the device buffer is insufficient. As a result, the tiling data is overwritten, the tiling mode is not limited by the tiling data copy speed and kernel execution speed, and the performance is further optimized.

  Figure 2 Tiling data moving with the kernel
  

Graph Operator Execution Process

The ATB uses an ordered Operation flow to process complex graph structures. You can call ATB APIs to build and optimize computational graphs.

Before using ATB, you need to sort out the entire computational graph and its computational flow, convert the topology of the computational graph into the FIFO structure, and construct the graph parameters of the ATB based on the sorted FIFO queue. For details about the conversion between the computational graph and FIFO queue, see Figure 3.

After obtaining the ordered Operation flow, the ATB traverses each Operation in the queue and executes the corresponding single-operator execution process.

Figure 3 Conversion between computational graphs and FIFO queues

Performance Optimization Mechanism

According to the actual data, the host process in the entire graph construction process takes a long time. To reduce the time consumed on the host and further optimize the end-to-end performance, the ATB provides four mechanisms for optimizing the host performance: Tiling Cache, Setup reusing, InferShape reusing, and Runner Pool.

• Tiling Cache

  The Tiling Cache is used to cache the tiling data of the kernel.

  According to the characteristics of the Transformer model, a large amount of kernel tiling data can be reused during inference. Therefore, the ATB caches the calculated tiling data. When reusable tiling data is detected, it is directly obtained from the cache instead of repeated calculation. In this way, a large amount of tiling data computation time is saved.

  As shown in Figure 4, the ATB contains two types of Tiling Caches: local Tiling Cache and global Tiling Cache (local cache and global cache for short).

  • The local cache is stored in the Operation object and can be read or written only by the current Operation object.
  • The global cache is a thread-level object. All Operation objects in the same thread can be read or written. Different from the local cache, the global cache is the sum of a group of caches.

  The number of caches contained in a global cache is determined by the number of Operations supported by the ATB.

  Figure 4 Two types of Tiling Caches in the ATB
  

Memory Optimization and Management Mechanism

The memory optimization and management mechanism mainly involves the calculation and management mechanism of the device memory by the ATB.

The device memory space required for delivering an operator consists of three parts: intermediate tensor memory, kernel scratch memory, and tiling data memory. The following describes how to calculate and manage the three parts in the ATB.

• The intermediate tensor is not used as the input or output of the entire kernel graph. Its lifecycle starts only when the first kernel that uses the intermediate tensor as the output tensor is executed. After all kernels that depend on the intermediate tensor as the input end, their lifecycle ends immediately.

  Therefore, at any given time, it is known which intermediate tensors exist in the entire kernel execution queue. The entire kernel execution process is essentially a process of continuously releasing and allocating intermediate tensors. The maximum amount of device memory occupied during this process can be considered as the memory required for intermediate tensors.

• The kernel scratch memory is used to store temporary data when the kernel is running.

  The ATB runs only one kernel on the same stream at a time. Therefore, different kernels can reuse the same scratch memory. The size of the scratch memory can be the maximum value among the capacities required by all kernels.

• The tiling data memory is the one that occupies less memory. The ATB allocates the total memory space of tiling data of all kernels at a time.

Figure 6 Calculation and management of the scratch memory and tiling data of the kernel

In the preceding three parts of memory, the intermediate tensor memory and the scratch memory of the kernel are allocated by the user as workspaces, and the tiling data of the kernel is managed and allocated by the ATB.

• The ATB calculates the intermediate tensor memory and scratch memory required by the entire graph in the Setup API, adds the two values, and returns the sum as WorkspaceSize to the user.
• The user allocates the device memory based on the WorkspaceSize returned by the Setup API, and then transfers the device memory to the ATB through the Execute API. The ATB uses the memory based on the previous calculation result.
• The kernel tiling data is stored in the context class of the ATB. When the context class is constructed, a 32 × 3 MB device memory pool is generated by default. Each time an Operation needs to move tiling data, a 3 MB device memory is fetched from the context to store tiling data. The number of memory blocks in the memory pool can be configured using environment variables. Currently, the size of each memory block cannot be configured.

Figure 7 ATB memory source

Context Introduction

The Context class is used to store and manage various public resources in the ATB. It contains the following resources: two streams, events for controlling time sequences, host memory pool, device memory pool, Runner pool, and overflow/underflow detection tensor.

• The two streams are used for kernel execution and copy of tiling data, respectively. The stream executed by the kernel is set by the user, and the stream for tiling data copy created by the ATB.

  If the multi-stream function is disabled, the stream used for tiling data copy will not be created.

• Events for controlling time sequences are used to ensure that the sequence of tiling data copy and kernel execution in the multi-stream function is correct.
• The host memory pool is a ring host memory buffer used to store tiling data computed on the host. The number of blocks can be controlled by ATB_HOST_TILING_BUFFER_BLOCK_NUM in Environment Variables.
• The structure and usage of the device memory pool are similar to those of the host memory pool. The difference is that the memory blocks in the device memory pool are device memory blocks, and the number of blocks can be controlled by ATB_DEVICE_TILING_BUFFER_BLOCK_NUM in Environment Variables.
• For details about the Runner pool, see "Runner Pool" in "Performance Optimization Mechanism". The Runner pool is stored in the context and is used by the Operation object as a public resource.
• The overflow/Underflow detection tensor is used to store the output result of the overflow/underflow detection operator. The device memory is allocated or freed with the Context.