PyTorch vector library optimization method, system and PyTorch vector library application method for RISC-V

Abstract

The present invention discloses a RISC-V-oriented PyTorch vector library optimization method, system, and PyTorch vector library application method, belonging to the field of computer software technology. The RISC-V-oriented PyTorch vector library optimization method includes: determining the RISC-V vector bit width supported by the host machine; based on the RISC-V vector bit width, enabling the corresponding vector extension compilation option to add the RVV dispatch scheduling mechanism to the PyTorch framework and encapsulate the RVV vector data type, thereby implementing a RISC-V-oriented PyTorch vector library. The present invention utilizes RISC-V vectors to extend the PyTorch vector library, thereby improving data processing efficiency in deep learning and machine learning applications.

Description

Technical Field

[0001] The present invention belongs to the field of computer software technology, and specifically proposes a PyTorch vector library optimization method and system for RISC-V, and a PyTorch vector library application method. Background Art

[0002] PyTorch, an open-source machine learning library developed by the Facebook AI research team, has been widely adopted in cutting-edge fields such as computer vision and natural language processing since its release in 2016. Its ease of use, flexibility, and powerful functionality have made PyTorch a popular choice in the research community. As an extension of Python, PyTorch provides a Python-based front-end interface that greatly simplifies the building and training of deep learning models. Among its core strengths, PyTorch's automatic differentiation system makes building deep neural networks more intuitive and efficient.

[0003] In the fields of machine learning and deep learning, the speed and efficiency of data processing are key to improving performance. To this end, modern processors widely adopt SIMD (Single Instruction, Multiple Data) technology, which allows a single instruction to perform operations on multiple data simultaneously, significantly improving data processing capabilities. PyTorch, a popular open source machine learning framework, has integrated support for multiple SIMD instruction sets, including AVX2, AVX512, NEON, SVE, and others, to fully utilize the parallel processing capabilities of modern processors.

[0004] While PyTorch already supports the various SIMD instruction sets mentioned above, it currently lacks support for RVV (RISC-V Vector Extension). RVV is a vector extension to the RISC-V architecture, designed to provide RISC-V processors with SIMD capabilities comparable to other architectures. As an open-source instruction set architecture, RISC-V has garnered widespread attention for its openness and customizability. The introduction of RVV heralds greater potential for RISC-V processors in fields such as machine learning and deep learning.

[0005] RISC-V is rapidly gaining momentum. Since its public debut in 2010, it has garnered support from numerous companies and organizations, including Intel, Apple, Google, ARM, and Huawei, and has been widely adopted in fields such as mobile devices, the Internet of Things, and artificial intelligence. With increasing global investment in processor technology, the RISC-V instruction set architecture holds a promising future. RISC-V is expected to play an increasingly important role in computer technology.

[0006] With the increasing popularity of the RISC-V architecture, more and more hardware platforms are beginning to support RVV. This presents new opportunities for PyTorch. By adding support for RVV, PyTorch can not only further expand its application scope on different hardware platforms, but also achieve significant performance improvements on these platforms.

[0007] In order to fully utilize the characteristics of the RISC-V vector instruction set, improve the performance and efficiency of PyTorch's basic operations and operators, allow the PyTorch vector library to fully utilize the advantages of SIMD instructions in the RISC-V instruction set, and significantly improve the running performance of PyTorch on the RISC-V platform, it is crucial to develop a new method to accelerate the PyTorch vector library. Summary of the Invention

[0008] In response to the technical problems existing in the prior art, the purpose of the present invention is to propose a PyTorch vector library optimization method, system and PyTorch vector library application method for RISC-V, which is used to optimize the PyTorch vector library using RISC-V Vector Extension (RVV) to improve the data processing efficiency in deep learning and machine learning applications.

[0009] To achieve the above object, the present invention adopts the following technical solutions:

[0010] A PyTorch vector library optimization method for RISC-V, the method comprising:

[0011] Determine the RISC-V vector bit width supported by the host;

[0012] Based on the RISC-V vector bit width, enable the corresponding vector extension compilation option to add the RVV dispatch scheduling mechanism to the PyTorch framework and encapsulate the RVV vector data type to implement the PyTorch vector library for RISC-V.

[0013] Furthermore, determining the RISC-V vector bit width supported by the host machine includes:

[0014] In the cmake file, run the vector program to determine whether the host hardware supports vector extensions;

[0015] If the host machine's hardware supports vector extensions, the vlenb built-in function is called to query the contents of the read-only control and status registers to obtain the RISC-V vector bit width supported by the host machine.

[0016] Furthermore, the vector extension compilation options include:

[0017] The target architecture is specified as RISC-V architecture, and supports general instruction set, compressed instruction set and vector extension instruction set;

[0018] and,

[0019] Specifies the length of the vector register.

[0020] Furthermore, the newly added RVV dispatch scheduling mechanism in the PyTorch framework includes:

[0021] Based on the environment detection mechanism, check the vector extension support of the host CPU;

[0022] If the host CPU supports vector extensions, initialize the RVV dispatch mechanism;

[0023] Register the RVV dispatch mechanism in PyTorch's kernel registration system to make it part of operator selection.

[0024] Furthermore, registering the RVV dispatch scheduling mechanism in the PyTorch kernel registration system includes:

[0025] Added REGISTER_RVV_DISPATCH macro definition, which is used to register PyTorch operators;

[0026] and,

[0027] Register the RVV operator in the specified operator, where the specified operators include: SparseBinaryOpIntersectionKernel operator, FlattenIndicesKernel operator, and BatchLinearAlgebraKernel operator.

[0028] Furthermore, encapsulating the RVV vector data type in the PyTorch framework includes:

[0029] Added rvv directory;

[0030] Encapsulate the RVV vector data type based on the RISC-V vector bit width, and control the RVV vector bit width to 256 bits through vl;

[0031] The encapsulated RVV vector data type is stored in the rvv directory.

[0032] A PyTorch vector library application method for RISC-V, the method comprising:

[0033] In the cmake file, run the vector program to determine whether the host hardware supports vector extensions;

[0034] When the current host machine supports vector extensions, the optimal data processing path is dynamically selected according to the data type of the current operation to implement basic operations on the encapsulated data type; wherein the PyTorch vector library is constructed based on any of the above-mentioned PyTorch vector library optimization methods for RISC-V, and the basic operations include: quantization operations, mathematical function operations, comparison operations, and memory access operations.

[0035] Furthermore, when the encapsulated data type is a float type, a qint8 type, a quint8 type, or a qint32 type, the basic operations on the encapsulated data type are implemented as follows:

[0036] Overload the class corresponding to the encapsulated RVV vector data type, and declare private variables based on the encapsulated RVV vector data type to store vector data;

[0037] In the case where the basic operation is a memory access operation,

[0038] If the package data type is float type and the address is aligned, a single-precision unit stride storage built-in function is used to store float type vector data;

[0039] If the encapsulated data type is qint32 type and the address is aligned, a 32-bit integer unit step storage built-in function is used to store qint32 type vector data;

[0040] If the encapsulated data type is float type or qint32 type, and the address is not aligned, the memory copy function memcpy function is called to store the float type or qint32 type vector data;

[0041] If the encapsulated data type is qint8 type, the 8-bit integer unit step storage built-in function is called to store the qint8 type vector data;

[0042] If the encapsulated data type is quint8 type, then the 8-bit unsigned integer unit step storage built-in function is called to store the quint8 type vector data;

[0043] In the case where the basic operation is a mathematical function operation, if the sleef vector math library supports it, the sleef library is called, otherwise scalar processing is performed and then merged into a vector;

[0044] When the basic operation is a comparison operation, a single-precision floating-point equality comparison built-in function is used to implement a vector equality operation, an unsigned integer negation built-in function is used to implement a vector inequality operation, a single-precision floating-point greater than comparison built-in function is used to implement a vector greater than operation, and a single-precision floating-point less than comparison built-in function is used to implement a vector less than operation.

[0045] Furthermore, in the case where the encapsulated data type is of type bf16, the basic operations on the encapsulated data type are implemented, including:

[0046] Convert bf16 type instructions to float type instructions;

[0047] Implement the corresponding basic operation according to the float type instruction to obtain the float type operation result;

[0048] Convert the float type operation result to the bf16 type operation result.

[0049] A PyTorch vector library optimization system for RISC-V, the system comprising:

[0050] A bit width determination module is used to determine the RISC-V vector bit width supported by the host;

[0051] A vector library implementation module that enables the corresponding vector extension compilation options based on the RISC-V vector bit width. This module adds the RVV dispatch scheduling mechanism to the PyTorch framework and encapsulates the RVV vector data type, thus implementing a PyTorch vector library for RISC-V.

[0052] Compared with the prior art, the present invention has the following positive effects:

[0053] (1) This paper specifically optimizes the PyTorch library's relatively weak performance on RISC-V and the problem that the advantages of the RISC-V vector instruction set are not fully utilized. It uses RVV intrinsics overloading to implement basic operations of different data types in the PyTorch vector library and is compatible with the existing PyTorch framework, ultimately improving the execution efficiency of the PyTorch API.

[0054] (2) The vector library of the present invention supports multiple data types, including float (single-precision floating point), bf16 (Brain Floating Point16), qint8 (quantized 8-bit signed integer), and quint8 (quantized 8-bit unsigned integer), to accommodate diverse computational requirements. Through optimization using RVV intrinsics, the present invention enables efficient processing of operations such as quantization, mathematical functions, comparisons, and memory accesses, significantly improving computational speed and performance. BRIEF DESCRIPTION OF THE DRAWINGS

[0055] Figure 1 It is a system architecture diagram of the present invention.

[0056] Figure 2 It is a diagram of the structure of the encapsulated float type. DETAILED DESCRIPTION

[0057] In order to make the various technical features and advantages or technical effects of the above technical solutions of the present invention more obvious and easy to understand, they are described in detail below with reference to the accompanying drawings.

[0058] Figure 1 This is a system architecture diagram of a PyTorch vector library optimization method for RISC-V and a PyTorch vector library application method proposed in the present invention. The steps of the PyTorch vector library optimization method for RISC-V are described in detail as follows.

[0059] Step 1: Determine the RISC-V vector width supported by the host.

[0060] Due to the scalability of the RISC-V instruction set, functions optimized for vector instructions cannot run on hardware that does not support vector extensions. Therefore, in the cmake file, the present invention determines whether the hardware supports vector extensions by running a vector program. If the hardware supports vector extensions, it further calls the vlenb built-in function __riscv_vlenb() to query the content of the read-only control and status register vlenb, that is, the vector register bit width VLEN_BITS.

[0061] Step 2: Based on the RISC-V vector bit width, enable the corresponding vector extension compilation option to add the RVV dispatch scheduling mechanism to the PyTorch framework and encapsulate the RVV vector data type.

[0062] 1) Enable vector extension compilation option.

[0063] Enable compilation options based on bit width. These compilation options include: specifying the target architecture as 64-bit RISC-V architecture, and supporting general instruction set, compressed instruction set, and vector extension instruction set, i.e. -march=rv64gcv; specifying the length of vector registers, i.e. -mrvv-vector-bits=${VLEN_BITS}.

[0064] 2) Added RVV dispatch mechanism

[0065] The RVV dispatch mechanism has been added to the PyTorch framework to dynamically select the optimal data processing path based on the vector scalability of the host machine. It includes the following sub-steps:

[0066] a) Environment detection and initialization: When PyTorch starts, the runtime environment detection mechanism checks the host CPU's support for vector extensions (RVV). If RVV support is detected, the RVV dispatch mechanism is initialized, including the priority of loading the PyTorch vector library to prepare for subsequent operator registration. If RVV is not supported, the default data processing path is used to ensure compatibility. Some code is shown below:

[0067]

[0068] b) RVV dispatch mechanism registration: Register the RVV dispatch mechanism in the PyTorch kernel registration system to make it part of the operator selection. Add the REGISTER_RVV_DISPATCH macro definition for registering PyTorch operators and register the RVV implementation branch in operators such as SparseBinaryOpIntersectionKernel, FlattenIndicesKernel, and BatchLinearAlgebraKernel. The following is the implementation of registering the RVV operator in SparseBinaryOpIntersectionKernel:

[0069]

[0070] 3) PyTorch vector library encapsulates multiple data types.

[0071] This paper adopts a modular design, encapsulating vector operations for different data types in independent modules for easier management and expansion. A new rvv directory is added to the PyTorch vector library to store encapsulations of different data types, maintaining compatibility with the existing PyTorch framework and ensuring seamless migration and use of existing code and models. The files encapsulating different data types are: vec256_bfloat16_rvv.h, vec256_float_rvv.h, vec256_qint8_rvv.h, vec256_quint8_rvv.h, and vec256_qint32_rvv.h.

[0072] Use the vector register bit width VLEN_BITS read in step 1 to encapsulate the RVV vector data type, and use vl to control the vector bit width to 256 bits. Taking the float type as an example, its encapsulation is as follows:

[0073]

[0074] In addition, the present invention also discloses a PyTorch vector library application method for RISC-V, which includes the following steps.

[0075] Step A: In the cmake file, run the vector program to determine whether the host hardware supports vector extensions.

[0076] When PyTorch executes an operation, the RVV dispatch mechanism dynamically selects the optimal data processing path based on the data type of the current operation and the RVV support of the host machine. When the user calls the PyTorch upper-level API interface, it will check whether the current host machine supports vector extensions RVV.

[0077] Step B: When the current host machine supports vector extension, dynamically select the optimal data processing path according to the data type of the current operation to implement basic operations on the encapsulated data type.

[0078] If the current host machine supports vector extensions (RVV), the PyTorch vector library implemented using RVV is selected for the registered operator, which improves the operator's running efficiency. Otherwise, the operator follows the default data processing path, that is, the scalar implementation. The basic operations implemented using RVV intrinsics include:

[0079] a) Quantization operation, which is used to convert high-precision data types into low-precision data types to reduce computing resource consumption;

[0080] b) Mathematical function operations, including trigonometric functions, exponential functions, logarithmic functions, etc.;

[0081] c) Comparison operations, used to perform comparisons of data types, such as ==, !=, <, >, etc.;

[0082] d) Memory access operation, used to handle memory access operations with unaligned addresses.

[0083] Specifically, if Figure 2 , using the RISC-V vector extension (RVV) built-in function intrinsics to implement basic operations of this data type.

[0084] In one embodiment, when the encapsulated data type is a float type, a qint8 type, a quint8 type, or a qint32 type, the steps of implementing basic operations of the data type include the following.

[0085] 1) Overload Vectorized <float>Class, and use the RVV vector data type encapsulated in step 3 to declare private variables as single-precision vector data types fixed_vfloat32m2_t values, which are used to store float type vector data; overload Vectorized <c10::qint8>Class, declare private variables as 8-bit integer vector data type fixed_vint8m2_tvalues, used to store qint8 type vector data; overload Vectorized <c10::quint8>Class, declare private variables as unsigned 8-bit integer vector data type fixed_vuint8m2_t values, used to store quint8 type vector data; overload Vectorized <c10::qint32>Class, declare private variables as 32-bit integer vector data type fixed_vint32m2_t values, used to store qint32 type vector data;

[0086] 2) Use the RVV built-in function intrinsics to implement the initialization of the overloaded class in 1), including the parameterless constructor and the parameterized constructor;

[0087] 3) Implement support for vector unaligned memory access. If the address is 4-byte aligned, use the single-precision unit step storage built-in function __riscv_vse32_v_f32m2 to store the float type, and use the 32-bit integer unit step storage built-in function __riscv_vse32_v_i32m2 to store the int32 type. Otherwise, call the memory copy function memcpy function for storage; for qint8 and quint8 types, there will be no problem of unaligned memory access. Then directly call the 8-bit integer unit step storage built-in function __riscv_vse8_v_i8m2 and the 8-bit unsigned integer unit step storage built-in function __riscv_vse8_v_u8m2 to store the qint8 type and quint8 type respectively;

[0088] 4) For mathematical operation functions, if the sleef vector math library supports it, the sleef library is called; otherwise, scalar processing is performed and then merged into vectors. For example, the sin function vector implementation calls Sleef_sinfx_u10rvvm2, the cos function vector implementation calls Sleef_cosfx_u10rvvm2, and the tan function vector implementation calls Sleef_tanfx_u10rvvm2.

[0089] 5) For comparison operations, the single-precision floating-point equal comparison built-in function __riscv_vmfeq_vv_f32m2_b16 and the unsigned integer negation built-in function __riscv_vnot_v_u32m2 are used to implement the vector operations of "==" and "!=", the single-precision floating-point greater than comparison built-in function __riscv_vmfgt_vv_f32m2_b16 is used to implement the vector operation of ">", and the single-precision floating-point less than comparison built-in function __riscv_vmflt_vv_f32m2_b16 is used to implement the vector operation of "<";

[0090] 6) For other basic operations, the corresponding RVV built-in functions intrinsics are also used to implement them.

[0091] In another embodiment, when the encapsulated data type is a float type, a qint8 type, a quint8 type, or a qint32 type, the basic operation for implementing the data type is to first convert a bf16 type instruction into a float type instruction, then implement the corresponding basic operation according to the float type instruction to obtain a float type operation result, and finally convert the float type operation result into a bf16 type operation result.

[0092] Among them, the conversion from bf16 to float data type is to convert Vectorized <bfloat16>Type data is converted into two Vectorized <float>Specifically, define a constant K, which is equal to Vectorized <bfloat16>The number of elements, that is, the number of bf16 types that can be contained in a vector, and then declare two arrays, arr is used to store the converted float type data, and arr2 is used to temporarily store BFloat16 type data. <bfloat16>The data of the type is stored in the arr2 array, and the PyTorch built-in conversion function is called to convert the BFloat16 data in arr2 into float type data and store it in arr;

[0093] To convert float to bf16 data type, two Vectorized <float>Types of data are merged and converted into a Vectorized <bfloat16>Specifically, define a constant K, which is equal to Vectorized <bfloat16>The number of elements, that is, the number of bf16 elements that can be contained in a vector, declares two arrays, arr is used to store the merged float type data, and arr2 is used to temporarily store the converted bf16 type data. Call the PyTorch built-in conversion function to convert the float data in arr to bf16 type data and store it in arr2.

[0094] We tested the performance of the Torch API using RVV intrinsics to implement the PyTorch vector library, and found that it far outperformed both scalar and auto-vectorization. The detailed data is shown in Table 1.

[0095] Table 1. Torch API performance comparison (unit: s)

[0096] scalar auto vectorization RVV intrinsics torch.sin 11.411 11.173 4.825 torch.cos 10.179 9.9 5.226 torch.tan 19.929 20.737 6.928 torch.round 8.726 7.435 2.024 torch.sqrt 3.211 1.188 1.137 torch.pow 3.115 0.368 0.36 torch.sum 1.953 0.305 0.268 torch.frac 4.801 5.721 0.734 torch.mul 4.423 0.408 0.411 torch.div 4.549 1.273 1.212

[0097] Although the present invention has been disclosed as above by way of embodiments, they are not intended to limit the present invention. Any appropriate modification or equivalent substitution of the technical solution of the present invention by a person skilled in the art should be included in the protection scope of the present invention. The protection scope of the present invention shall be based on that defined in the claims. < / float> < / float> < / float>

Claims

1. A PyTorch vector library optimization method for RISC-V, characterized in that: The method comprises: Determine the RISC-V vector bit width supported by the host; Based on the RISC-V vector bit width, enable the corresponding vector extension compilation option to add the RVVdispatch scheduling mechanism to the PyTorch framework and encapsulate the RVV vector data type to implement the PyTorch vector library for RISC-V; The newly added RVV dispatch scheduling mechanism in the PyTorch framework includes: Based on the environment detection mechanism, check the vector extension support of the host CPU; If the host CPU supports vector extensions, initialize the RVV dispatch mechanism; Register the RVV dispatch mechanism in PyTorch's kernel registration system to make it part of operator selection; Registering the RVV dispatch mechanism in the PyTorch kernel registration system includes: Added REGISTER_RVV_DISPATCH macro definition, which is used to register PyTorch operators; and, Register the RVV operator in the specified operator, which includes: SparseBinaryOpIntersectionKernel operator, FlattenIndicesKernel operator, and BatchLinearAlgebraKernel operator; The encapsulation of the RVV vector data type in the PyTorch framework includes: Added rvv directory; Encapsulate the RVV vector data type based on the RISC-V vector bit width, and control the RVV vector bit width to 256 bits through vl; The encapsulated RVV vector data type is stored in the rvv directory.

2. The RISC-V-oriented PyTorch vector library optimization method according to claim 1, characterized in that: Determining the RISC-V vector bit width supported by the host includes: In the cmake file, run the vector program to determine whether the host hardware supports vector extensions; If the host machine's hardware supports vector extensions, the vlenb built-in function is called to query the contents of the read-only control and status registers to obtain the RISC-V vector bit width supported by the host machine.

3. The RISC-V-oriented PyTorch vector library optimization method according to claim 1, characterized in that: The vector extension compilation options include: The target architecture is specified as RISC-V architecture, and supports general instruction set, compressed instruction set and vector extension instruction set; and, Specifies the length of the vector register.

4. A PyTorch vector library application method for RISC-V, characterized in that: The method comprises: In the cmake file, run the vector program to determine whether the host hardware supports vector extensions; When the current host machine supports vector extensions, the optimal data processing path is dynamically selected according to the data type of the current operation to implement basic operations on the encapsulated data type; wherein the PyTorch vector library is constructed based on the PyTorch vector library optimization method for RISC-V according to any one of claims 1 to 3, and the basic operations include: quantization operations, mathematical function operations, comparison operations, and memory access operations.

5. The RISC-V-oriented PyTorch vector library application method according to claim 4, characterized in that: When the encapsulated data type is a float type, a qint8 type, a quint8 type, or a qint32 type, the basic operations on the encapsulated data type are implemented as follows: Overload the class corresponding to the encapsulated RVV vector data type, and declare private variables based on the encapsulated RVV vector data type to store vector data; In the case where the basic operation is a memory access operation, If the package data type is float type and the address is aligned, a single-precision unit stride storage built-in function is used to store float type vector data; If the encapsulated data type is qint32 type and the address is aligned, a 32-bit integer unit step storage built-in function is used to store qint32 type vector data; If the encapsulated data type is float type or qint32 type, and the address is not aligned, the memory copy function memcpy function is called to store the float type or qint32 type vector data; If the encapsulated data type is qint8 type, the 8-bit integer unit step storage built-in function is called to store the qint8 type vector data; If the encapsulated data type is quint8 type, then the 8-bit unsigned integer unit step storage built-in function is called to store the quint8 type vector data; In the case where the basic operation is a mathematical function operation, if the sleef vector math library supports it, the sleef library is called, otherwise scalar processing is performed and then merged into a vector; When the basic operation is a comparison operation, a single-precision floating-point equality comparison built-in function is used to implement a vector equality operation, an unsigned integer negation built-in function is used to implement a vector inequality operation, a single-precision floating-point greater than comparison built-in function is used to implement a vector greater than operation, and a single-precision floating-point less than comparison built-in function is used to implement a vector less than operation.

6. The RISC-V-oriented PyTorch vector library application method according to claim 4, characterized in that: In the case where the encapsulated data type is of type bf16, the basic operations on the encapsulated data type are implemented as follows: Convert bf16 type instructions to float type instructions; Implement the corresponding basic operation according to the float type instruction to obtain the float type operation result; Convert the float type operation result to the bf16 type operation result.

7. A PyTorch vector library optimization system for RISC-V, characterized in that: The system comprises: A bit width determination module is used to determine the RISC-V vector bit width supported by the host; A vector library implementation module, which is used to enable the corresponding vector extension compilation option based on the RISC-V vector bit width, to implement the PyTorch vector library for RISC-V after adding the RVV dispatch scheduling mechanism and encapsulating the RVV vector data type in the PyTorch framework; The newly added RVV dispatch scheduling mechanism in the PyTorch framework includes: Based on the environment detection mechanism, check the vector extension support of the host CPU; If the host CPU supports vector extensions, initialize the RVV dispatch mechanism; Register the RVV dispatch mechanism in PyTorch's kernel registration system to make it part of operator selection; Registering the RVV dispatch mechanism in the PyTorch kernel registration system includes: Added REGISTER_RVV_DISPATCH macro definition, which is used to register PyTorch operators; and, Register the RVV operator in the specified operator, which includes: SparseBinaryOpIntersectionKernel operator, FlattenIndicesKernel operator, and BatchLinearAlgebraKernel operator; The encapsulation of the RVV vector data type in the PyTorch framework includes: Added rvv directory; Encapsulate the RVV vector data type based on the RISC-V vector bit width, and control the RVV vector bit width to 256 bits through vl; The encapsulated RVV vector data type is stored in the rvv directory.

Images