An ai accelerator optimization method and system based on a multi-path kernel software stack

Abstract

The application discloses an AI accelerator optimization method and system based on a multi-path kernel software stack, characterized by adopting zero-copy, automatic reverse address mapping and dynamic load balancing methods, realizing high-compatibility accelerator access support, realizing efficient inference through a double-path architecture, providing wide compatibility through a user-mode daemon calling a user-mode accelerator library in a slow path, realizing kernel direct access to an accelerator through hijacking a working environment in a fast path, and introducing zero-copy data transmission, automatic reverse address mapping, dynamic load balancing and an adaptive water level mechanism to further optimize memory management and task scheduling. Compared with the prior art, the application has the advantages of reducing inherent delay and end-to-end delay of kernel machine learning inference, eliminating kernel-user space communication overhead, and significantly improving system performance, and solves the problems of high inference delay and poor compatibility caused by the fact that an operating system kernel cannot directly access a hardware accelerator, and has a good application prospect.

Description

Technical Field

[0001] This invention relates to the field of operating system and machine learning fusion technology, specifically to an AI accelerator optimization method and system based on a multi-path kernel software stack. Background Technology

[0002] In user space, the interaction between a process and an accelerator typically occurs as follows: The process first calls the accelerator library's initialization interface and shares a segment of memory allocated on the host (called the working environment) and a register located in the accelerator (called the doorbell). Subsequent calls to the accelerator are packaged into a specific format and placed in this shared memory. The accelerator is then notified via hardware to check this shared memory, retrieve the task within, and execute it. However, accelerator vendors only provide user-space accelerator libraries and do not support kernel-space access. With the development of emerging scenarios such as autonomous driving, augmented reality, and intelligent robots, the operating system kernel needs to integrate machine learning models to quickly optimize the performance of its components under new loads, such as I / O prediction, memory management, and load balancing. However, current operating system kernels cannot directly access hardware accelerators. Common solutions include: 1) making remote calls through a user-space daemon, which introduces significant kernel-user space communication overhead, leading to a substantial increase in inference latency; 2) reimplementing the entire accelerator library in kernel space, which requires significant manpower for maintenance and makes it difficult to guarantee compatibility and stability.

[0003] In summary, existing operating system kernels suffer from high inference latency and poor compatibility due to their inability to directly access hardware accelerators. Therefore, providing efficient and highly compatible accelerator access support within the kernel is a pressing issue that needs to be addressed. Summary of the Invention

[0004] The purpose of this invention is to address the shortcomings of existing technologies by providing an AI accelerator optimization method and system based on a multi-path kernel software stack. It employs a slow-path and fast-path architecture in the operating system kernel, optimizing memory management and task scheduling to achieve efficient inference and improve system performance. The slow path utilizes a user-space daemon process to call the user-space accelerator library for broad compatibility, while the fast path hijacks the working environment to enable direct kernel access to the accelerator, eliminating kernel-user space communication overhead. Furthermore, this invention introduces zero-copy data transfer, automatic reverse address mapping, dynamic load balancing, and adaptive watermark mechanisms to further optimize memory management and task scheduling. This reduces the inherent latency of kernel machine learning inference to 10.37% of existing technologies and end-to-end latency to 17.33%, significantly improving system performance and demonstrating promising application prospects and commercial development value.

[0005] The specific technical solution to achieve the purpose of this invention is: an AI accelerator optimization method based on a multi-path kernel software stack, characterized by the use of zero-copy, automatic reverse address mapping, and dynamic load balancing methods to optimize memory management and task scheduling, achieving highly compatible accelerator access support. This method specifically includes the following steps:

[0006] Step 1: Path Division

[0007] Paths are divided according to function type: initialization functions take the slow path, and inference functions take the fast path.

[0008] Step 2: Initialization

[0009] The user-space daemon calls the accelerator library initialization interface to establish the device context and working environment.

[0010] Step 3: Work Environment Hijacking

[0011] Map the user-mode working environment to the kernel address space to enable direct kernel access;

[0012] Step 4: Reasoning

[0013] The kernel hijacks the working environment and writes it directly into the task descriptor, then uses a doorbell mechanism to notify the accelerator to execute.

[0014] The path segmentation step determines the type of accelerator interface to be called. For environment initialization interfaces, it jumps to the initialization step; for inference interfaces, it jumps to the inference step. The initialization step creates a daemon process in user space. This daemon process can be pre-created at system startup or when the invention starts. The parameters required for the initialization interface of the user-space accelerator library are then sent to the daemon process, which calls the initialization interface provided by the user-space accelerator library and returns the result to the kernel. The working environment hijacking step obtains the working environment (memory allocated in the host) from the result returned by the initialization step and maps it to kernel space. Typically, the result returned by the initialization step is a user-space address, which can be resolved once or multiple times to find the memory address of the target working environment, and then mapped to the kernel address space. The inference step stores the data from the called interface in the working environment in a format consistent with the user-space library and notifies the accelerator to retrieve it via a doorbell (a register in the accelerator). During the initialization phase, the inference step uses user-space forwarding to call the accelerator's zero-copy memory interface to allocate host memory shared by the daemon process and the accelerator, and transmits the user address to the kernel and maps it to the kernel address space.

[0015] Before sending computation tasks through the inference step, the zero-copy step requires the user to store data in the accelerator's memory, effectively mitigating the overhead of transferring data to the accelerator via the daemon process. Specifically, during initialization, the accelerator's zero-copy memory interface is called via user-mode forwarding to allocate host memory shared by the daemon process and the accelerator. Data stored in this memory is automatically retrieved by the accelerator. Subsequently, the user address is transferred to the kernel and mapped to the kernel address space. Before inference, the data is written to this memory via the kernel address, and the inference interface is called with the user-mode address as a parameter, placing it in the working environment for direct access by the accelerator. As a complement to the zero-copy step, this invention provides an automatic reverse mapping step. The mapping relationship is recorded when mapping user-mode addresses to kernel addresses, allowing the user to use the kernel-mode address when calling the inference interface. The interface automatically searches the mapping relationship to find the corresponding user-mode address and replaces the kernel-mode address as the actual parameter.

[0016] The zero-copy method allocates device-visible memory through the user-mode library during the initialization phase and maps it to the kernel space, enabling the kernel to directly write data. The automatic reverse address mapping records the relationship between kernel-mode and user-mode address mapping, so that the kernel can automatically convert kernel addresses to user-mode addresses when calling the inference interface. The dynamic load balancing method utilizes the function allocation engine to delay task distribution through a buffered scheduling queue and dynamically selects accelerators or CPUs for execution based on device utilization and predicted latency.

[0017] The function allocation engine includes an adaptive water level mechanism and configuration for dynamically managing the number of tasks in the accelerator's internal queue. The adaptive water level mechanism includes low, medium, and high water levels, and dynamically adjusts the medium and high water levels based on the number of remaining tasks when tasks are replenished, in order to balance equipment utilization and the risk of task backlog.

[0018] The inference step includes a load balancing step, ensuring that inference tasks are not immediately sent to the accelerator's working environment. Instead, they remain in an additional buffer queue, and tasks are retrieved from this queue only when the number of tasks in the accelerator's working environment is low. In addition to the accelerator, the CPU is also considered a device that can retrieve tasks from the queue for execution.

[0019] The load balancing step provides an adaptive water level mechanism to dynamically manage the number of tasks within the accelerator's operating environment. This adaptive water level mechanism includes low, medium, and high water levels. When the number of tasks in the operating environment reaches the medium water level, tasks are added, with the number being the difference between the medium and high water levels. When the added tasks arrive in the operating environment, the remaining number of tasks is set as the low water level, and the high water level is adjusted upwards, so that the medium water level is at the average position of the high and low water levels.

[0020] An operating system based on a multi-path kernel software stack for AI accelerator optimization is characterized by employing a slow-path and fast-path architecture. This operating system comprises kernel-mode and user-mode components. The kernel-mode component can be compiled with the system or dynamically loaded as a kernel module to implement the fast path, directly accessing the accelerator's working environment and submitting inference tasks. The user-mode component can be created at startup or during initialization to implement the slow path, using a daemon process to call the user-mode accelerator library for initialization and support of non-critical functions.

[0021] Compared with existing technologies, this invention significantly improves system performance. It initializes the working environment using a user-mode library through user-mode forwarding to ensure compatibility with a large number of accelerators supported by the user-mode library. Then, it hijacks the working environment to kernel mode. During inference, the hijacked working environment communicates directly with the accelerator, thereby enabling efficient inference. It also reduces the inherent latency of kernel machine learning inference to 10.37% of existing technologies and the end-to-end latency to 17.33%, showing good application prospects and commercial development value. Attached Figure Description

[0022] Figure 1 This is a system diagram of the architecture of the present invention;

[0023] Figure 2 Diagram showing accelerator access;

[0024] Figure 3 This is a schematic diagram of an adaptive water level mechanism. Detailed Implementation

[0025] See Figure 1 This invention employs an operating system with slow and fast path architectures. The operating system includes a kernel-mode part and a user-mode part, which implements AI accelerator optimization based on a multi-path kernel software stack. The kernel-mode part can be compiled with the system or dynamically loaded as a kernel module to implement the fast path, directly accessing the accelerator working environment and submitting inference tasks. The user-mode part can be created at startup or during initialization to implement the slow path, using a daemon process to call the user-mode accelerator library to complete initialization and support non-critical functions.

[0026] The AI ​​accelerator optimization method based on a multi-path kernel software stack of the present invention specifically includes:

[0027] Step 1: Path Division

[0028] Paths are divided according to function type: initialization functions take the slow path, and inference functions take the fast path.

[0029] Step 2: Initialization

[0030] The user-space daemon calls the accelerator library initialization interface to establish the device context and working environment.

[0031] Step 3: Work Environment Hijacking

[0032] Map the user-mode working environment to the kernel address space to enable direct kernel access;

[0033] Step 4: Reasoning

[0034] The kernel hijacks the working environment and writes it directly into the task descriptor, then uses a doorbell mechanism to notify the accelerator to execute.

[0035] See Figure 2 The path segmentation step determines the type of accelerator interface to be called. For the environment initialization interface, it jumps to the initialization step; for the inference interface, it jumps to the inference step. The initialization step creates a daemon process in user space. This daemon process can be pre-created at system startup or when the invention starts. Then, the parameters required for the initialization interface of the user-space accelerator library are sent to the daemon process, which calls the initialization interface provided by the user-space accelerator library and returns the result to the kernel. The working environment hijacking step obtains the working environment (memory allocated in the host) from the result returned by the initialization step and maps it to kernel space. Usually, the result returned by the initialization step is a user-space address. After one or more address resolutions, the memory address of the target working environment can be found and then mapped to the kernel address space. The inference step stores the data of the called interface in the working environment in a format consistent with the user-space library and notifies the accelerator to retrieve it via a doorbell (a register in the accelerator). During the initialization phase, the inference step uses user-space forwarding to call the accelerator's zero-copy memory interface to allocate host memory shared by the daemon process and the accelerator, and transmits the user address to the kernel and maps it to the kernel address space.

[0036] Before sending computation tasks through the inference step, the zero-copy step requires the user to store data in the accelerator's memory, effectively mitigating the overhead of transferring data to the accelerator via the daemon process. Specifically, during initialization, the accelerator's zero-copy memory interface is called via user-mode forwarding to allocate host memory shared by the daemon process and the accelerator. Data stored in this memory is automatically retrieved by the accelerator. Subsequently, the user address is transferred to the kernel and mapped to the kernel address space. Before inference, the data is written to this memory via the kernel address, and the inference interface is called with the user-mode address as a parameter, placing it in the working environment for direct access by the accelerator. As a complement to the zero-copy step, this invention provides an automatic reverse mapping step. The mapping relationship is recorded when mapping user-mode addresses to kernel addresses, allowing the user to use the kernel-mode address when calling the inference interface. The interface automatically searches the mapping relationship to find the corresponding user-mode address and replaces the kernel-mode address as the actual parameter.

[0037] The zero-copy method allocates device-visible memory through the user-mode library during the initialization phase and maps it to the kernel space, enabling the kernel to directly write data. The automatic reverse address mapping records the relationship between kernel-mode and user-mode address mapping, so that the kernel can automatically convert kernel addresses to user-mode addresses when calling the inference interface. The dynamic load balancing method utilizes the function allocation engine to delay task distribution through a buffered scheduling queue and dynamically selects accelerators or CPUs for execution based on device utilization and predicted latency.

[0038] The function allocation engine includes an adaptive water level mechanism and configuration for dynamically managing the number of tasks in the accelerator's internal queue. The adaptive water level mechanism includes low, medium, and high water levels, and dynamically adjusts the medium and high water levels based on the number of remaining tasks when tasks are replenished, in order to balance equipment utilization and the risk of task backlog.

[0039] The inference step includes a load balancing step, ensuring that inference tasks are not immediately sent to the accelerator's working environment. Instead, they remain in an additional buffer queue, and tasks are retrieved from this queue only when the number of tasks in the accelerator's working environment is low. In addition to the accelerator, the CPU is also considered a device that can retrieve tasks from the queue for execution.

[0040] See Figure 3 The load balancing step provides an adaptive water level mechanism to dynamically manage the number of tasks within the accelerator's working environment. This adaptive water level mechanism includes low, medium, and high water levels. When the number of tasks in the working environment reaches the medium water level, tasks are added, with the number being the difference between the medium and high water levels. When the added tasks arrive in the working environment, the remaining number of tasks in the working environment is set as the low water level, and the high water level is adjusted upwards, so that the medium water level is at the average position of the high and low water levels.

[0041] To better understand the solutions of this invention, the invention will be further described in detail below with reference to specific embodiments. Obviously, the described embodiments are merely some, not all, of the embodiments of this invention. All other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this invention.

[0042] Example 1

[0043] The following is an example of providing support for AMD's ROCm accelerator library. This example is implemented as a kernel module and includes the following steps:

[0044] Step 1: After startup, the kernel loads the kernel module implemented in the example and starts the user-mode daemon process.

[0045] Step 2: The daemon calls the ROCm initialization interface has_init, then calls hsa_queue_create to create the working environment and doorbell and returns the result.

[0046] Step 3: The kernel module sets the address offset according to the device model to find the working environment and doorbell address. It then maps the working environment address to the kernel space using vmap and the doorbell address to the kernel space using ioremap, thus completing the working environment hijacking.

[0047] Step 4: When the kernel module receives a memory allocation task, it increases the required space size by 32 bytes and packages it as a parameter. It then communicates with the daemon process via the netlink mechanism to send the parameter to the daemon process. The daemon process calls hsa_amd_memory_pool_allocate to allocate memory and sends it back to the kernel module in preparation for the zero-copy step.

[0048] Step 5: The kernel module maps the allocated memory to kernel space using vmap, saves a fixed magic number and user space address at the starting position, then manages the mapping range through a mapple tree for reverse mapping, and finally returns the starting position plus 32 bytes as the final result.

[0049] Step 6: The kernel module directly reads and writes data through the kernel address and transmits it directly to the accelerator using zero-copy.

[0050] Step 7: When the kernel module receives the inference task, it will first look up the 32 bytes preceding the user-provided kernel address to see if they contain the magic number. If they do, it will convert the user-space address to the corresponding user-space address through the recorded user-space address. Otherwise, it will look up the corresponding user-space address in the maple tree and then put the parameters into the automatic load balancing buffer queue according to the device requirements.

[0051] Step 8: Automatic load balancing will work automatically, placing tasks into the device's working environment and setting the doorbell.

[0052] The above is merely a further description of the present invention and is not intended to limit the scope of this patent. Any equivalent implementation of the present invention should be included within the scope of the claims of this patent.

Claims

1. An AI accelerator optimization method based on a multi-path kernel software stack, characterized in that, By employing zero-copy, automatic reverse address mapping, and dynamic load balancing methods, memory management and task scheduling are optimized to achieve highly compatible accelerator access support. This method specifically includes the following steps: Step 1: Path Division Paths are divided according to function type: initialization functions take the slow path, and inference functions take the fast path. Step 2: Initialization The user-space daemon calls the accelerator library initialization interface to establish the device context and working environment. Step 3: Work Environment Hijacking Map the user-mode working environment to the kernel address space to enable direct kernel access; Step 4: Reasoning The kernel directly writes the task descriptor through the hijacked working environment and notifies the accelerator to execute via a doorbell mechanism; the zero-copy allocates device-visible memory through the user-mode library during the initialization phase and maps it to the kernel space, enabling the kernel to directly write data; the automatic reverse address mapping records the relationship between kernel-mode and user-mode address mapping, so that the kernel automatically converts kernel addresses to user-mode addresses when calling the inference interface; the dynamic load balancing utilizes a function allocation engine to delay task distribution through a buffered scheduling queue and dynamically selects an accelerator or CPU for execution based on device utilization and predicted latency.

2. The AI ​​accelerator optimization method based on a multi-path kernel software stack according to claim 1, characterized in that, The function allocation engine includes an adaptive water level mechanism and configuration for dynamically managing the number of tasks in the accelerator's internal queue. The adaptive water level mechanism includes low, medium, and high water levels, and dynamically adjusts the medium and high water levels based on the number of remaining tasks when tasks are replenished, in order to balance equipment utilization and the risk of task backlog.

3. An operating system based on the AI ​​accelerator optimization method architecture according to claim 1, characterized in that, An operating system employing a slow-path and fast-path architecture comprises two parts: kernel mode and user mode. The kernel mode implements the fast path, directly accessing the accelerator's working environment and submitting inference tasks. The user mode implements the slow path, using a daemon process to call the user-mode accelerator library to complete initialization and support non-critical functions.

Images