TPU-based tensor extension calculation method and device, medium and product

Abstract

The invention discloses a tensor extension calculation method and device based on TPU, a medium and a product, and relates to the technical field of chips, and the method comprises the steps: analyzing a tensor extension request, and obtaining the first dimension information of an original tensor and the second dimension information of a target tensor; determining a tensor expansion demand according to the first dimension information and the second dimension information; according to the tensor expansion requirement, data, needing to be subjected to dimension expansion, in the original tensor is carried to a control memory or a vector memory from a high-bandwidth memory of the TPU, and in combination with a tensor broadcasting mechanism on the TPU, target tensor data is formed and stored in the high-bandwidth memory again through parallel processing of all the core computing units. According to the technical scheme of the embodiment of the invention, a tensor extension algorithm capable of accurately adapting to TPU hardware features is provided, so that the TPU native supports the extension calculation of the tensor dimension, the blank of the TPU function is filled, and the calculation parallelism and the transmission efficiency in the tensor extension calculation are improved.

Description

Technical Field

[0001] This invention relates to the field of computer underlying technology, and more particularly to the field of chip technology, specifically to a tensor expansion computing method, device, medium, and product based on a TPU (Tensor Processing Unit). Background Technology

[0002] Current deep learning model training and inference heavily rely on tensor operations (such as matrix multiplication and convolution). However, traditional general-purpose processors, such as CPUs (Central Processing Units) or GPUs (Graphics Processing Units), face problems such as low energy efficiency and insufficient memory bandwidth when processing these tasks. To address this, the TPU (Telematics Processing Unit) was developed. Through its dedicated systolic array and MXU (Matrix Unit) design, the TPU can significantly improve the parallel efficiency of neural network computation.

[0003] Currently, various neural network computations often require frequent expansion of tensor dimensions, such as feature copying or dimension alignment. However, existing TPU hardware does not natively support tensor expansion operations. In current technologies, tensor expansion logic is only implemented on CPUs and GPUs using a "view" mechanism. If tensor expansion needs to be forcibly implemented on a TPU, it is necessary to iterate through each data element and expand the tensor by copying each element. This implementation method has extremely low processing efficiency, huge hardware overhead, and cannot effectively utilize the various hardware advantages inherent in TPUs. Summary of the Invention

[0004] This invention provides a tensor expansion computation method, device, medium, and product based on TPU, to provide a tensor expansion algorithm that can accurately adapt to the hardware characteristics of TPU, enabling TPU to natively support tensor dimension expansion computation.

[0005] According to one aspect of the present invention, a tensor expansion computation method based on a TPU is provided. The TPU includes two parallel core computing units and a hierarchical memory system, the hierarchical memory system including high-bandwidth memory, control memory, and vector memory. The method includes:

[0006] The tensor expansion request is parsed to obtain the first dimension information of the original tensor and the second dimension information of the target tensor.

[0007] Based on the first and second dimension information, determine the tensor expansion requirements;

[0008] Based on the tensor expansion requirements, the data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or vector memory. Combined with the tensor broadcasting mechanism on the TPU, the target tensor data is formed and stored back in the high-bandwidth memory through parallel processing of each core computing unit.

[0009] According to another aspect of the present invention, a tensor expansion computing device based on a TPU is provided. The TPU includes two parallel core computing units and a hierarchical memory system, the hierarchical memory system including high-bandwidth memory, control memory, and vector memory. The device includes:

[0010] The dimension information acquisition module is used to parse the tensor expansion request and obtain the first dimension information of the original tensor and the second dimension information of the target tensor.

[0011] The extension requirement determination module is used to determine the tensor extension requirements based on the first dimension information and the second dimension information.

[0012] The extended processing module is used to move the data in the original tensor that needs to be expanded in dimension from the high-bandwidth memory of the TPU to the control memory or vector memory according to the tensor expansion requirements. Combined with the tensor broadcasting mechanism on the TPU, the target tensor data is formed and stored back in the high-bandwidth memory through parallel processing of each core computing unit.

[0013] According to another aspect of the present invention, an electronic device is also provided, the electronic device comprising:

[0014] At least one processor; and

[0015] A memory communicatively connected to the at least one processor; wherein,

[0016] The memory stores a computer program that can be executed by the at least one processor, which enables the at least one processor to perform the TPU-based tensor expansion computation method according to any embodiment of the present invention.

[0017] According to another aspect of the present invention, a computer-readable storage medium is provided, the computer-readable storage medium storing computer instructions, the computer instructions being configured to cause a processor to execute and implement the TPU-based tensor expansion computation method described in any embodiment of the present invention.

[0018] According to another aspect of the present invention, a computer program product is also provided, including a computer program that, when executed by a processor, implements the steps of the TPU-based tensor expansion computation method as described in any embodiment of the present invention.

[0019] The technical solution of this invention parses the tensor expansion request to obtain the first dimension information of the original tensor and the second dimension information of the target tensor; based on the first and second dimension information, the tensor expansion requirement is determined; according to the tensor expansion requirement, the data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or vector memory, and combined with the tensor broadcast mechanism on the TPU, the target tensor data is re-stored in the high-bandwidth memory through parallel processing of each core computing unit. This method can adaptively utilize the hardware features of the TPU, such as hierarchical memory and dual-core computing units, to perform tensor dimension expansion calculations according to the actual tensor expansion requirements. It can solve the problems of inefficient memory access, low hardware resource utilization, and excessive data movement under large data volumes caused by directly porting the tensor expansion logic of CPU or GPU to the TPU. It creatively proposes a tensor expansion algorithm that can accurately adapt to the hardware features of the TPU, enabling the TPU to natively support tensor dimension expansion calculations, effectively filling the gap in TPU functionality, and significantly improving the computational parallelism and transmission efficiency in tensor expansion calculations.

[0020] It should be understood that the description in this section is not intended to identify key or essential features of the embodiments of the present invention, nor is it intended to limit the scope of the invention. Other features of the invention will become readily apparent from the following description. Attached Figure Description

[0021] To more clearly illustrate the technical solutions in the embodiments of the present invention, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0022] Figure 1 This is a flowchart of a TPU-based tensor expansion computation method provided according to an embodiment of the present invention;

[0023] Figure 2 This is a flowchart of another TPU-based tensor expansion computation method provided by an embodiment of the present invention;

[0024] Figure 3 This is a schematic diagram of the structure of a TPU-based tensor expansion computing device according to an embodiment of the present invention;

[0025] Figure 4 This is a schematic diagram of the structure of an electronic device that implements the TPU-based tensor expansion computation method according to an embodiment of the present invention. Detailed Implementation

[0026] To enable those skilled in the art to better understand the present invention, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort should fall within the scope of protection of the present invention.

[0027] It should be noted that the terms "first," "second," etc., in the specification, claims, and accompanying drawings of this invention are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of the invention described herein can be implemented in orders other than those illustrated or described herein. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover non-exclusive inclusion; for example, a process, method, system, product, or apparatus that comprises a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or apparatus.

[0028] Figure 1 This is a flowchart of a TPU-based tensor expansion computation method provided in an embodiment of the present invention. This embodiment is applicable to situations where various hardware resources in a TPU are effectively utilized to efficiently implement tensor expansion computation (typically, the expand operator). This method can be executed by a TPU-based tensor expansion computation device, which can be implemented in hardware and / or software and is generally configured in a TPU. For example, it can be executed by the instruction controller in the TPU, or it can be configured in other types of terminal devices to cooperate with the TPU for execution.

[0029] The TPU specifically comprises two parallel core computing units (also known as TensorCores) and a hierarchical memory system. The hierarchical memory system includes High Bandwidth Memory (HBM), Control Memory (CMEM), and Vector Memory (VMEM). Each core computing unit contains a Vector Unit (VPN) for performing general element-wise operations and multiple MXUs. Each core computing unit also contains a Vector Memory module.

[0030] Correspondingly, such as Figure 1 As shown, the method may include:

[0031] S110. Parse the tensor expansion request to obtain the first dimension information of the original tensor and the second dimension information of the target tensor.

[0032] The tensor expansion request refers to an expansion instruction used to expand one or more dimensions of the original tensor to form the target tensor. The tensor expansion request can be sent from the software layer to the instruction controller of the TPU.

[0033] Typically, the tensor expansion request can be generated by making operator calls to tensor expansion operators (e.g., the expand operator) in the software stack.

[0034] Tensor expansion is a fundamental operation in deep learning, used to expand a tensor from size 1 to a target size (or keep the original size) while preserving the data content. This operation is essential in features alignment or batch processing in neural networks, especially in scenarios with large data volumes where memory transfer efficiency is extremely important. In other words, tensor expansion can only be performed on dimensions of size 1.

[0035] By performing semantic parsing on the tensor expansion request, we can obtain the dimensional information of the original tensor (i.e., the first dimensional information) and the dimensional information of the target tensor (i.e., the second dimensional information).

[0036] In a specific example, for a tensor expansion request that expands an original tensor tensor1 with tensor size (1,2,4,1) to a target tensor tensor2 with tensor size (128,2,4,256), (1,2,4,1) is the first dimension information and (128,2,4,256) is the second dimension information.

[0037] Furthermore, by parsing the tensor expansion request, in addition to obtaining the first and second dimension information, we can also obtain the first storage address of the original tensor and the second storage address of the target tensor. Both the first and second storage addresses are storage addresses on high-bandwidth memory. Using the first storage address, the original tensor stored in high-bandwidth memory can be retrieved and then moved to control memory or vector memory in one go or in batches for tensor expansion processing by the core computing unit. Simultaneously, the tensor expansion result can also be moved back to the second storage address on high-bandwidth memory in one go or in batches and stored as the target tensor.

[0038] S120. Determine the tensor expansion requirements based on the first-dimensional information and the second-dimensional information.

[0039] Tensor expansion requirements can be understood as the information regarding which dimensions of the original tensor need to be expanded. Specifically, these requirements may include at least one expanded dimension and the corresponding expansion dimensionality for each dimension. In this embodiment, the inventors, through creative analysis of the characteristics of each level of memory in the hierarchical memory system of the TPU and the parallel computing mechanism of the two core computing units, have categorized tensor expansion requirements into three types of dimensionality expansion requirements that may occur simultaneously. One type is the requirement to expand the lowest dimension of the original tensor; another type is the requirement to expand the non-lowest dimensions to a small to medium scale; and the third type is the requirement to expand the non-lowest dimensions to a large scale.

[0040] In a specific example, assuming that the tensor expansion request can expand the tensor dimension of a five-dimensional tensor to a maximum of five dimensions, then the aforementioned minimum dimension can refer to the 0th dimension, and the aforementioned non-minimum dimensions can refer to the 1st to 4th dimensions.

[0041] Optionally, small-to-medium dimensionality and large-scale dimensionality can be distinguished based on a preset dimensionality threshold (e.g., 256). In a specific example, when dimensional expansion within 256 dimensions is required, it can be defined as small-to-medium dimensionality expansion; when expansion beyond 256 dimensions is required, it can be defined as large-scale dimensionality expansion. The preset dimensionality threshold is determined based on the adaptability of memory capacity at each level of the hierarchical memory system.

[0042] Continuing the previous example, assuming the rightmost dimension in the dimension information is the lowest dimension, when the first dimension information is (1,2,4,1) and the second dimension information is (128,2,4,256), the tensor expansion requirement includes: the requirement to expand the lowest dimension (0th dimension) by 256, and the requirement to expand the third dimension by a small to medium scale of 128.

[0043] S130. Based on the tensor expansion requirements, the data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or vector memory. Combined with the tensor broadcasting mechanism on the TPU, the target tensor data is formed and stored in the high-bandwidth memory through parallel processing of each core computing unit.

[0044] In this embodiment, depending on one or more types of expansion requirements included in the tensor expansion needs, the data requiring dimensional expansion in the original tensor can be moved from the TPU's high-bandwidth memory to the control memory or vector memory. One or both core computing units can then directly use the original tensor stored in the control memory or vector memory to perform the tensor expansion calculation. Specifically, the required tensor expansion task can be evenly divided, and two core computing units can execute half of the tensor expansion task in parallel to improve the utilization efficiency of the TPU's computing power.

[0045] For each core computing unit, multiple MXUs within the core computing unit can be used to perform matching vector parallel operations on the assigned tensor expansion task, with a single instruction capable of processing multiple data elements. Specifically, each core computing unit can directly call various instructions from the TPU's dedicated instruction set, such as vector loading instructions or data storage instructions, to execute the corresponding tensor expansion task based on its internal MXUs.

[0046] In this embodiment, since the MXUs in the TPU can directly apply the hardware broadcast mechanism to implicitly expand the dimension of small tensors based on the pulsating array. Furthermore, by using multiple MXUs in parallel to execute tensor expansion tasks in each core computing unit, data transfer overhead can be effectively reduced and energy efficiency can be improved.

[0047] In other words, in this embodiment, by using different methods to coordinate the collaborative processing of various hardware components in the TPU for different tensor expansion needs, the rich MXUs on the TPU and the contents of different storage capacities in the hierarchical memory system can be fully utilized to efficiently realize tensor expansion calculations in a collaborative manner with dual-core computing units.

[0048] The technical solution of this invention parses the tensor expansion request to obtain the first dimension information of the original tensor and the second dimension information of the target tensor; based on the first and second dimension information, the tensor expansion requirement is determined; according to the tensor expansion requirement, the data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or vector memory, and combined with the tensor broadcast mechanism on the TPU, the target tensor data is re-stored in the high-bandwidth memory through parallel processing of each core computing unit. This method can adaptively utilize the hardware features of the TPU, such as hierarchical memory and dual-core computing units, to perform tensor dimension expansion calculations according to the actual tensor expansion requirements. It can solve the problems of inefficient memory access, low hardware resource utilization, and excessive data movement under large data volumes caused by directly porting the tensor expansion logic of CPU or GPU to the TPU. It creatively proposes a tensor expansion algorithm that can accurately adapt to the hardware features of the TPU, enabling the TPU to natively support tensor dimension expansion calculations, effectively filling the gap in TPU functionality, and significantly improving the computational parallelism and transmission efficiency in tensor expansion calculations.

[0049] Figure 2 This is a flowchart of another TPU-based tensor expansion calculation method provided by an embodiment of the present invention. This embodiment is an optimization based on the above embodiments. In this embodiment, the operation of "according to the tensor expansion requirements, the data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or vector memory, and combined with the tensor broadcasting mechanism on the TPU, the target tensor data is formed and re-stored in the high-bandwidth memory through the parallel processing of each core computing unit" is specified.

[0050] Correspondingly, such as Figure 2 As shown, the method may include:

[0051] S210. Parse the tensor expansion request to obtain the first dimension information of the original tensor and the second dimension information of the target tensor.

[0052] S220. Based on the first-dimensional information and the second-dimensional information, determine the tensor expansion requirements.

[0053] S230. Based on the data access granularity of the TPU, perform memory alignment processing on the lowest dimension data of the original tensor.

[0054] In this embodiment, considering that the core computing unit in the TPU hardware (typically the MXU in the core computing unit) needs to load multiple data in a single cycle in fixed-size blocks (e.g., 128 bytes) when accessing memory, if the original tensor does not achieve the aforementioned memory alignment in high-bandwidth memory, additional shifting and merging operations will be added. Therefore, before performing subsequent tensor expansion operations, memory alignment of the lowest dimension data of the original tensor is first required. Because the higher-dimensional data in the tensor is always an integer multiple of the lowest-dimensional data, memory alignment only needs to be performed on the lowest-dimensional data. For example, the ALIGN128() function can be used directly to adjust the starting address of the memory block for the lowest-dimensional data in the original tensor to ensure that data is stored at an address aligned to 128 bytes.

[0055] S240. Divide the vector memory into an input area and an output area, wherein the input area is used to store data in the original tensor and the output area is used to store data in the target tensor.

[0056] In this embodiment, considering that the general data processing flow based on TPU is to first move the data to be processed from the high-bandwidth memory to the vector memory inside the core computing unit through the DMA (Direct Memory Access) controller, after the core computing unit directly obtains the data from the vector memory and performs the calculation, it temporarily stores the calculation result data in the vector memory, and finally moves the calculation result data back from the vector memory to the high-bandwidth memory through the DMA controller.

[0057] That is, the vector memory needs to store both the input data to be processed and the calculation result data to be output. In order to effectively avoid data conflicts, in this embodiment, the vector memory is first divided into an input area and an output area before performing tensor expansion calculation.

[0058] Optionally, to facilitate convenient management of the input and output areas, the input and output areas can be divided equally. For example, assuming the size of the vector memory on the TPU is 3072*1024 bytes, the first 1536*1024 bytes of the vector memory can be used as the input area, and the last 1536*1024 bytes of the vector memory can be used as the output area, and the starting addresses of the input and output areas can be set accordingly.

[0059] S250. Following the order from low to high dimensions, sequentially obtain a target expansion dimension and the corresponding target expansion dimension from the tensor expansion requirements, and then execute the following branch process:

[0060] If the target expanded dimension is the lowest dimension, execute S260; if the target expanded dimension is not the lowest dimension and the target expanded dimension is less than or equal to the preset dimension threshold, execute S270; if the target expanded dimension is not the lowest dimension and the target expanded dimension is greater than the preset dimension threshold, execute S280.

[0061] Specifically, when expanding the dimensions of a tensor, the expansion proceeds from lower dimensions to higher dimensions. Furthermore, following the order from lower dimensions to higher dimensions (typically, from dimension 0 -> dimension 1 -> dimension 2 -> dimension 3 -> dimension 4), a target expansion dimension and its corresponding target expansion number can be sequentially obtained from the tensor expansion requirements.

[0062] Furthermore, based on the target expansion dimension and its specific value, different tensor expansion methods are selected for tensor expansion processing to maximize the efficient utilization of TPU hardware resources.

[0063] S260: Move the first data matching the target extended dimension from high-bandwidth memory to vector memory. Combined with the tensor broadcast mechanism on the TPU, and through the parallel processing of each core computing unit, after generating the first extended data matching the target extended dimension in vector memory, move the first extended data back to high-bandwidth memory and then execute S290.

[0064] In an optional implementation of this embodiment, when the target expansion dimension is the lowest dimension, the first data matching the target expansion dimension is moved from high-bandwidth memory to vector memory. Combined with the tensor broadcast mechanism on the TPU, and through parallel processing of each core computing unit, after generating the first expanded data matching the target expansion dimension in vector memory, the first expanded data is moved back to high-bandwidth memory. Specifically, this may include:

[0065] S2601. Detect whether the amount of data formed when expanding the first data to the first expanded data can be completely stored in the vector memory.

[0066] In this optional embodiment, the inventors consider that the amount of data expanded is minimal when performing the lowest-dimensional data expansion. Therefore, it is possible that the first data can be moved to the vector memory all at once for expansion calculation, and the resulting expansion calculation result (i.e., the first expanded data) can be completely stored in the vector memory. In this case, there is no need for block-based data transfer; the DMA controller only needs to transfer the data once, resulting in the highest data transfer efficiency. Therefore, it is necessary to first calculate whether the amount of data formed when expanding the first data to the first expanded data can be completely stored in the vector memory.

[0067] S2602. If the data can be stored completely, the first data is moved from the high-bandwidth memory to the input area of ​​the vector memory in one go, and the data expansion task matching the target expansion dimension is split into two sub-tasks and distributed to two core computing units, and S2603 is executed.

[0068] In this embodiment, a data expansion task is first generated for expanding the first data to the first expanded data, targeting the expansion dimension. Then, this data expansion task is split into two equal groups, which are then distributed to two core computing units in the TPU (typically device0 and device1). This setup enables parallel processing of data expansion, effectively improving overall data throughput.

[0069] S2603 uses two core computing units and a tensor broadcast mechanism to call TPU vector instructions in parallel to perform data copying operations on the assigned subtasks, forming the first extended data stored in the output area of ​​the vector memory, and then moving the first extended data completely stored in the output area back to the high-bandwidth memory in one go.

[0070] In this embodiment, when constructing the data expansion task, it is necessary to first calculate the step size of the input tensor and the step size of the output tensor. This step size can be understood as the address increment of the dimension expansion. In addition, in order to use the tensor broadcasting mechanism, it is also necessary to determine the target expanded dimension as the number of broadcasts, and then mark the step size of the broadcast dimension as 0.

[0071] Accordingly, in this branch, the step size of dimension 0 needs to be marked as 0. This setting ensures data reuse when expanding the broadcast dimension; that is, data is reused when the broadcast dimension address remains unchanged, effectively reducing the frequent use of the DMA controller for data transfer.

[0072] In this system, data expansion is achieved through batch processing of vector instructions in vector memory, and the broadcast dimension is implemented through vector copying. The broadcast task can be calculated separately by two core computing units, each representing (number of broadcasts) / 2.

[0073] Furthermore, for each core computing unit, eight internal parallel MXUs can be used to execute the corresponding data expansion (copying) tasks in parallel. Specifically, TPU vector instructions can be used directly within each core computing unit to achieve batch data processing. That is, after loading and copying a single 128-byte data element, a single instruction can process the batch loading and storage of 8*128 bytes of that data element at once, which can give full play to the parallel advantages of MXUs (vector processors or vector registers) and improve the utilization rate of TPU hardware resources.

[0074] S2604. If the data cannot be stored completely, determine the number of first blocks that match the first data, and move the first data to the input area of ​​the vector memory in batches according to the number of first blocks, and execute S2605.

[0075] In this embodiment, if the vector memory cannot fully store the data volume formed when the first data is expanded to the first expanded data, it is necessary to first determine the number of first blocks matching the first data based on the memory capacity of the vector memory. This number of first blocks can be understood as the need to split the first data into several blocks, and input these blocks in batches into the input area of ​​the vector memory for tensor expansion processing, ensuring that the data processed in a single batch does not exceed the capacity of the vector memory.

[0076] S2605. Split the data expansion task corresponding to the first data of each batch into two subtasks and distribute them to two core computing units, and execute S2606.

[0077] S2606 uses two core computing units and employs a tensor broadcast mechanism to call TPU vector instructions in parallel to perform data copying operations on the assigned subtasks, thereby forming the first extended data of each batch stored in the output area of ​​the vector memory, and then moving the first extended data of each batch back to the high-bandwidth memory in batches.

[0078] S270: Move the second data that matches the target expansion dimension from the high-bandwidth memory to the control memory. Combined with the tensor broadcast mechanism on the TPU, through the parallel processing of each core computing unit, directly store the second expanded data that matches the target expansion dimension obtained by expanding the second data back into the high-bandwidth memory, and then execute S290.

[0079] In an optional implementation of this embodiment, when the target expansion dimension is not the lowest dimension and the target expansion dimension is less than or equal to a preset dimension threshold, the second data matching the target expansion dimension is moved from high-bandwidth memory to control memory. Combined with the tensor broadcast mechanism on the TPU, and through parallel processing of each core computing unit, the second expanded data matching the target expansion dimension obtained by expanding the second data is directly re-stored into high-bandwidth memory. Specifically, this may include:

[0080] S2701: Move the second data matching the target expansion dimension from the high-bandwidth memory to the control memory in one go; split the data expansion task matching the target expansion dimension into two sub-tasks and distribute them to two core computing units.

[0081] In this embodiment, considering that the capacity of the control memory is much larger than that of the vector memory, it is generally twice that of the vector memory. The inventors calculated that when the target expansion dimension is not the lowest dimension and the target expansion dimension is less than or equal to a preset dimension threshold (e.g., 256 dimensions), the original tensor and the second expanded data obtained by tensor expansion processing of the original tensor can be stored directly in the control memory at one time.

[0082] At this point, the original tensor can be moved from the high-bandwidth memory to the control memory in one go through the DMA controller. After the two core computing units generate the second extended data based on the original tensor stored in the control memory, the second extended data can be moved back from the control memory to the high-bandwidth memory in one go through the DMA controller again.

[0083] In an optional implementation of this embodiment, a core computing unit, such as device0, can move the complete raw tensor from high-bandwidth memory to control memory via a DMA controller. After ensuring the entire transfer process is complete (e.g., by calling the dlc_sync() function) and ensuring inter-device synchronization between the two core computing units (e.g., by calling the sync_device() function), the two core computing units are then triggered to execute the tensor expansion task in parallel to prevent the second core computing unit, such as device1, from reading invalid data.

[0084] S2702: Through two core computing units, a tensor broadcast mechanism is used to call TPU vector instructions in parallel to perform data copying operations on the assigned subtasks, and the expanded second extended data is re-stored in high-bandwidth memory.

[0085] In this optional implementation, the target expansion dimension also needs to be used as the number of broadcasts, and the target expansion dimension is temporarily marked as 1 so as to facilitate the processing of the original data in a single block.

[0086] Specifically, each core computing unit can directly use the DMA controller to copy data from control memory to a storage address in high-bandwidth memory that matches the target tensor. This process is repeated until both core computing units have jointly completed all broadcast data copy operations, resulting in the final target tensor.

[0087] It should be noted that, since the second data is directly copied in the control memory, and the control memory can support any step size (i.e., no alignment is required), it is possible to flexibly meet the requirements for irregular expansion of non-zero dimensions.

[0088] S280: Move the third data that matches the target extended dimension from high-bandwidth memory to vector memory in blocks. Combined with the tensor broadcast mechanism on the TPU, and through the parallel processing of each core computing unit, after generating the third extended data that matches the target extended dimension in blocks in vector memory, move the third extended data back to high-bandwidth memory in blocks, and then execute S290.

[0089] In an optional implementation of this embodiment, when the target expansion dimension is not the lowest dimension and the target expansion dimension is greater than a preset dimension threshold, the third data matching the target expansion dimension is moved from high-bandwidth memory to vector memory in blocks. Combined with the tensor broadcast mechanism on the TPU, and through parallel processing of each core computing unit, after the third expansion data matching the target expansion dimension is generated in blocks in vector memory, the third expansion data is moved back to high-bandwidth memory in blocks. Specifically, this may include:

[0090] S2801. Determine the number of second blocks that match the third data, and move the third data to the input area of ​​the vector memory in batches according to the number of second blocks.

[0091] In this optional implementation, when the target expansion dimension is not the lowest dimension and the target expansion dimension is greater than the preset dimension threshold, considering that even if control memory is used, it is impossible to completely store the third data and the third expansion data that matches the third data, vector memory is still used to implement the above tensor expansion operation.

[0092] Specifically, it is necessary to first calculate the computational dimension within each block and the dimension of the block transmission to determine the processing scope of the third data in a single batch. Furthermore, it is also necessary to set the broadcast flag for the target expanded dimension to 0 and use the target expanded dimension as the total number of broadcasts.

[0093] S2802, split the data expansion task corresponding to the third data of each batch into two subtasks and distribute them to two core computing units.

[0094] S2803 uses two core computing units and employs a tensor broadcast mechanism to call TPU vector instructions in parallel to perform data copying operations on the assigned subtasks, thereby forming the third extended data of each batch stored in the output area of ​​the vector memory, and then moving the third extended data of each batch back to the high-bandwidth memory in batches.

[0095] S290. Determine whether the processing of all extended dimensions is complete: if yes, execute S2100; otherwise, return to execute S250.

[0096] S2100: Obtain the target tensor data stored in high-bandwidth memory.

[0097] The technical solutions of the embodiments of the present invention have the following technical advantages:

[0098] 1. Use memory alignment instructions in the TPU, such as ALIGN128() or ALIGN256(), to ensure that the data storage address matches the TPU vector access granularity and avoid the loss of access due to misalignment.

[0099] 2. By using dual-core computing units for parallel processing during tensor expansion operations, data throughput can be effectively improved.

[0100] 3. Dynamically calculate block size based on vector memory capacity to avoid memory overflow and effectively reduce DMA cycles.

[0101] 4. Broadcast Dimension DMA Multiplexing: Data is transmitted only once for broadcast dimensions (dimensions with a step size of 0), reducing the number of accesses to high-bandwidth memory.

[0102] 5. When performing small-scale tensor expansion calculations for non-zero dimensions (dimensions 1-4), the large capacity of the control memory (twice that of vector memory) is utilized to reduce the number of block and DMA data transfers under large data volumes. At the same time, the combination of H2C (host to storage device) and C2H (storage device to host) DMA provides flexible adaptability for diverse future step size requirements and enhances code scalability.

[0103] Furthermore, in order to describe the technical solutions of the various embodiments of the present invention in more detail, the tensor expansion calculation process of the embodiments of the present invention will now be described in detail.

[0104] First, the main computational flow of tensor expansion computation is as follows:

[0105] 1. Input a tensor expansion request that expands the original tensor to the target tensor.

[0106] 2. Initialize the tensor expansion request to obtain the identification information (device_id) of the two core computing units and the storage address of the original tensor on HBM, and parse the tensor dimension information (also known as shape information) of the original tensor and the target tensor.

[0107] 3. Record the original size of the lowest dimension of the original tensor (also known as dim0) and store it in the input_dim1_unpad variable. By saving the original size of dim0 before memory alignment, we can determine whether dim0 is a broadcast dimension later.

[0108] 4. Perform memory alignment operations on dim0. For example, the input and output dim0 can be processed using the ALIGN128() function.

[0109] 5. Determine whether extended computation is needed for higher dimensions (typically 1-4 dimensions) of the original tensor:

[0110] If a high-dimensional dim(x) is defined as x belonging to [1,2,3,4], no extended calculation is required, or the required extended output dimension is less than or equal to 256, then the broadcast flag (broadcast_flag) of dim(x) is set to 1, and the broadcast count (dma_times) and single DMA transfer length (dma_length) corresponding to dim(x) are calculated. Here, dma_length can be understood as the sum of all low-dimensional elements contained in dim(x).

[0111] If a high-dimensional dim(j) is required to perform extended calculations, and the required output dimension of the high-dimensional dim(j) is greater than 256, then the broadcast flag of the dim(x) is set to 0.

[0112] 6. For each dimension that has been processed, call the Broadcast function to perform the tensor expansion operation as described above, in order from lower dimensions to higher dimensions.

[0113] In the Broadcast function, three branches are used to perform corresponding tensor expansion operations for three different types of tensor expansion requirements. Before executing the above three branches, the following steps 1-5 are performed first.

[0114] 1. The VMEM needs to be partitioned, that is, the VMEM is divided into an input area and an output area, each of which is 1536*1024 bytes.

[0115] 2. Based on the dimension information of the target tensor, obtain the last non-1 dimension in sequence (which can be understood as obtaining one dimension to be expanded in sequence from low dimension to high dimension) as output_lastdim.

[0116] 3. Split the extended tasks executed for output_lastdim according to device_id. For example, distribute the load tasks on the first core computing unit XYS0 (i.e., the aforementioned device0) and the second core computing unit XYS1 (i.e., the aforementioned device1).

[0117] 4. Calculate the stride (stride_i) of the input tensor and the stride (stride_o) of the output tensor by calling the stride calculation function (Calculate_stride() function).

[0118] 5. First, the step size of the broadcast dimension of output_lastdim is set to 0 using the broadcast_judge() function to ensure data reuse during tensor expansion.

[0119] After completing operations 1-5 above, perform tensor expansion computation in any of the following three branches.

[0120] Branch 1: broadcast_flag=0, input_dim1_unpad=1, that is, the current dimension being processed is the lowest dimension, and the lowest dimension needs to be expanded.

[0121] At this point, `dim0_broadcast=1` is set, and the `Calculate_batchDim` function is used to calculate whether the total size of the data generated when performing tensor expansion calculations on the lowest dimension exceeds the VMEM. If so, the split size (`dim_split`) is calculated based on the total data size and the VMEM's memory capacity. The lowest dimension data is then divided into blocks according to `dim_split`, and multiple batches of data are processed cyclically. That is, the DMA controller moves the data blocks of the split size for each batch from HBM to the VMEM's input area. For broadcast data, the DMA controller only moves the data once. If not, the DMA controller loads the lowest dimension data stored in HBM into the VMEM's input area all at once. Then, for this broadcast dimension (the lowest dimension), a single data element is loaded at a time, and 128 copies are generated using vector copying to achieve single-element copying to a vector via broadcast. Finally, the copied vector is stored in the VMEM's output area. Finally, the DMA controller writes the replication results of each batch from the VMEM output area back to the HBM, or writes all the replication results from the VMEM output area back to the HBM at once.

[0122] Branch 2: When the required output dimension of output_lastdim is greater than 256, broadcast_flag=0. In this case, the split dimension (batch_dim) and split size (dim_split) need to be calculated accordingly. Then, the data of output_lastdim is divided into blocks according to dim_split, and multiple batches of data are processed in a loop. That is, the data blocks of the split size of each batch are moved from HBM to the input area of ​​VMEM for data processing through the DMA controller, and then moved back to HBM in blocks.

[0123] Branch 3: broadcast_flag=1, meaning the required output dimension for output_lastdim is less than or equal to 256. In this case, the required computational data is transferred from HBM to CMEM as a unified data source via DMA through one of the core computing units (e.g., device 0). After device synchronization between the two core computing units, the broadcast task can be further split according to device_id, i.e., XYS0 and XYS1 each handle half of the task. Then, the DMA controller continuously executes batch copy tasks from CMEM to HBM. After completing dma_times broadcast tasks, the CMEM broadcast ends, and global device synchronization is performed to confirm the consistency of the HBM output tensors.

[0124] Finally, the specific implementation process of each branch will be described in detail.

[0125] I. 0-Dimensional (dim0) Extension Branch

[0126] 1. Core objective: To handle the expansion or alignment requirements of the lowest dimension (dim0), adapt to the characteristics of the TPU vector computing unit (128-byte parallel granularity), and achieve efficient 1D to ND expansion or data alignment through VMEM, avoiding computational efficiency loss caused by dimension mismatch.

[0127] 2. Judgment Logic (Triggering Conditions)

[0128] (1) Basic conditions: The lowest dimension (dim0) of the tensor is the target dimension of the expansion operation. Due to the TPU hardware limitation, it needs to be aligned to a multiple of 128 bytes.

[0129] (2) Detailed judgment:

[0130] Detailed judgment 1: Whether expansion is needed: Determined by input_dim1_unpad (the original size before dim0 alignment):

[0131] <1> If input_dim1_unpad == 1 (equal to 1): it needs to be expanded from 1 dimension to N dimensions (e.g., output dim0 = 256), triggering the broadcast logic.

[0132] <2> If input_dim1_unpad != 1 (not equal to 1): only alignment is needed (e.g., input dim0=50 → output dim0=128), no broadcasting is required.

[0133] Detailed judgment 2: Whether block division is needed: Determined by the Calculate_batchDim function:

[0134] <1> If batch_dim=0 is returned, it means that the expanded data volume exceeds the capacity of a single VMEM region (1536×1024), and it needs to be divided into blocks.

[0135] <2> If batch_dim ≥ 1 is returned, it means that the amount of data is less than or equal to the VMEM capacity, and can be processed in one go.

[0136] 3. Processing flow

[0137] (1) For extended and aligned calculation scenarios:

[0138] <1> Scenario expansion is needed:

[0139] Step size marker: The broadcast_judge function sets stride_i[0] to 0 (indicating that no offset is needed when broadcasting).

[0140] Vector copy: Load one element into x0 using the inherent vector instruction v_f32_ld_tnsr_st_msk in the TPU, and then copy it to the entire vector register (1→128 parallel expansion) using x0 = x0[0].

[0141] <2> Alignment only for the scene:

[0142] Step size is preserved: stride_i[0] is the original calculated value (non-zero, based on the actual offset).

[0143] Batch loading: Data is loaded using the vector instruction v_f32_ld_tnsr_st_msk inherent in the TPU, and valid elements are marked with a mask.

[0144] (2) For data transmission and chunking scenarios:

[0145] <1> Segmentation scenario:

[0146] Splitting logic: Split dim0 by dim_split (e.g., 128), and increment dim_split in the loop variable dma_dim1 each time.

[0147] Transmission process: Each batch of data is transferred from HBM to VMEM via the DMA controller, processed, and then written back from VMEM to HBM until all blocks are completed.

[0148] <2> No block partitioning scenario:

[0149] One-time transfer: All data is transferred to VMEM through a single DMA operation, processed, and then written back to HBM.

[0150] II. Grouped CMEM Broadcast Branch

[0151] 1. Core objective: To handle small to medium-sized expansions (target size ≤ 256) with non-zero dimensions (1-4 dimensions), and to reduce the number of DMA controller startups by leveraging the large capacity and low access latency characteristics of CMEM, thereby achieving efficient batch broadcast replication.

[0152] 2. Judgment Logic (Triggering Conditions)

[0153] (1) Basic conditions: There is an expansion requirement for non-zero dimensions (1-4 dimensions) (input size ≠ output size), and the target expansion size is ≤256.

[0154] (2) Flag Trigger: When the above conditions are met, broadcast_flag is set to 1, the number of broadcasts dma_times (equal to the output dimension size) is recorded, and the dimension is temporarily set to 1 (to facilitate processing the original data in single blocks).

[0155] 3. Processing flow

[0156] (1) Data source preparation:

[0157] Device 0 (device_id=0) transmits raw data from HBM to CMEM via the DMA controller as a unified data source for broadcasting.

[0158] Call the dlc_sync() function to ensure the transmission is complete, and then use sync_device() to achieve synchronization between devices, preventing device 1 (device_id=1) from reading invalid data.

[0159] (2) Parallel broadcasting by two devices:

[0160] <1> Task splitting: dividing dma_times by device_id:

[0161] Device 0 processes the first dma_times / 2 broadcasts, with the offset starting from the address set for dma_length (incrementing dma_length each time); Device 1 processes the remaining broadcasts, with the offset starting from the beginning and end of (dma_times / 2+1)*dma_length.

[0162] <2> Batch copy: Copy CMEM data to the HBM target address via dlc_dma, repeating this process until all broadcasts are completed.

[0163] (3) Hardware adaptation:

[0164] CMEM's large capacity (twice that of VMEM) stores the complete data source, avoiding repeated readings from HBM by multiple devices and reducing HBM bandwidth usage. At the same time, it adapts to CMEM's support for arbitrary step sizes (no alignment required), flexibly addressing irregular expansion needs for non-zero dimensions.

[0165] III. General Block-Based VMEM Branch

[0166] 1. Core objective: To handle large-scale expansion of non-zero dimensions (1-4 dimensions) (target size > 256), adapt to VMEM capacity limitations through dynamic block division, and achieve efficient expansion in general scenarios by combining dual-core computing units (XYS0 / XYS1) in parallel.

[0167] 2. Judgment Logic (Triggering Conditions)

[0168] (1) Basic conditions: There is an expansion requirement for non-zero dimensions (1-4 dimensions), and the target size is >256 (which does not meet the capacity threshold of grouped CMEM).

[0169] (2) Blocking basis: The Calculate_batchDim() function traverses dimensions 1-4, finds the first dimension whose cumulative data volume exceeds the VMEM capacity as batch_dim (splitting dimension), and calculates dim_split (single block size, ensuring that single block ≤ VMEM capacity).

[0170] 3. Processing flow

[0171] (1) Dimensional splitting and parameter calculation:

[0172] <1> Calculate dimension parameters: Call the output_MathSzie() function to calculate the computation dimension within the block, and call the output_DmaSzie() function to calculate the dimension of the block transfer, thus defining the scope of a single batch processing.

[0173] <2> Step size adjustment: Set stride=0 for dimensions of size 1 using broadcast_judge() to enable broadcast marking.

[0174] (2) Chunked transmission and computation:

[0175] <1> Multi-dimensional block loop: Iterate through dimensions 1-4 using the output_DmaSzie() function, and split the batch_dim dimension by dim_split (e.g., dim2=1000 bytes → split into 8 blocks, each block is 125 bytes).

[0176] <2> Data reuse optimization: If the current block has the same address as the previous block (broadcast scenario), reuse VMEM data to avoid repeatedly performing DMA operations.

[0177] <3> Vector-level computation: In VMEM, the dimensions are traversed by output_MathSzie, and data is copied using vector instructions (since the extension is essentially a copy, no complex calculations are required).

[0178] (3) Dual-device parallel operation:

[0179] <1> The task is split based on output_lastdim (the last non-1 dimension) and device_id, with device 0 handling the first half and device 1 handling the second half.

[0180] <2> By setting up independent input and output areas in VMEM to isolate data regions, conflicts between devices can be avoided.

[0181] (4) Hardware adaptation:

[0182] Dynamic partitioning allows large-dimensional data to fit the VMEM capacity, avoiding memory overflow; parallel processing by dual computing units improves throughput and adapts to the characteristics of TPU multi-core architecture.

[0183] It should be emphasized again that the various embodiments of the present invention, through adaptive improvements to the tensor expansion calculation algorithm for TPU hardware, achieve the following technical effects:

[0184] 1. Improved computational efficiency

[0185] Compared to the serial element-by-element copying method of the CPU, the embodiments of the present invention utilize the vector computing unit of the TPU (such as an 8×128 vector register) to achieve batch data expansion. A single instruction can handle the parallel copying of multiple elements (the specific number matches the width of the vector register), reducing the number of loop iterations. At the same time, through task splitting of dual-core computing units (device_id=0, device_id=1), parallel processing of data expansion is achieved, further improving the overall throughput.

[0186] 2. Memory access optimization

[0187] To address the memory access characteristics of the TPU, alignment mechanisms such as the ALIGN128() function are used to ensure that data addresses meet the hardware access granularity requirements, avoiding instruction splitting overhead caused by unaligned accesses. At the same time, VMEM is used as an intermediate cache to reduce the frequency of access to HBM.

[0188] 3. Improved hardware resource utilization

[0189] Existing solutions do not fully utilize the vector computing capabilities and multi-core architecture of TPUs, and mostly process data in a scalar manner. The embodiments of this invention leverage the parallel advantages of vector registers through vector instructions (such as v_f32_ld_tnsr_st_msk) and achieve load balancing between the two core computing units based on device identifiers, thereby improving the utilization rate of TPU hardware resources.

[0190] 4. Optimization of Dimension Expansion Logic Adaptability

[0191] Existing solutions employ a uniform processing flow for all dimensional expansion scenarios, failing to differentiate between high and low dimensional characteristics. The embodiments of this invention demonstrate the design of differentiated processing logic for expansion requirements that are not the highest dimensional and are the highest dimensional. This reduces the number of DMA operations under large data volumes, while its architecture provides an adaptability basis for diverse requirements at future step sizes, enhancing code scalability.

[0192] 5. Data consistency guarantee

[0193] The embodiments of the present invention ensure the accuracy of data copying during the expansion process through strict address mapping (based on step size calculation) and synchronization mechanism. That is, the elements with dimension 1 in the original tensor are accurately expanded to the target size, and the expanded data content is completely consistent with the theoretical expectation, which fundamentally guarantees the correctness of the tensor operator calculation.

[0194] Figure 3 This is a schematic diagram of a tensor expansion computing device based on a TPU provided in an embodiment of the present invention. The TPU includes two parallel core computing units and a hierarchical memory system, which includes high-bandwidth memory, control memory, and vector memory, such as... Figure 3 As shown, the device includes:

[0195] The dimension information acquisition module 310 is used to parse the tensor expansion request and obtain the first dimension information of the original tensor and the second dimension information of the target tensor.

[0196] The extension requirement determination module 320 is used to determine the tensor extension requirements based on the first dimension information and the second dimension information.

[0197] The extended processing module 330 is used to move the data in the original tensor that needs to be expanded in dimension from the high-bandwidth memory of the TPU to the control memory or vector memory according to the tensor expansion requirements. Combined with the tensor broadcasting mechanism on the TPU, the target tensor data is formed and stored in the high-bandwidth memory through parallel processing of each core computing unit.

[0198] The technical solution of this invention parses the tensor expansion request to obtain the first dimension information of the original tensor and the second dimension information of the target tensor; based on the first and second dimension information, the tensor expansion requirement is determined; according to the tensor expansion requirement, the data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or vector memory, and combined with the tensor broadcast mechanism on the TPU, the target tensor data is re-stored in the high-bandwidth memory through parallel processing of each core computing unit. This method can adaptively utilize the hardware features of the TPU, such as hierarchical memory and dual-core computing units, to perform tensor dimension expansion calculations according to the actual tensor expansion requirements. It can solve the problems of inefficient memory access, low hardware resource utilization, and excessive data movement under large data volumes caused by directly porting the tensor expansion logic of CPU or GPU to the TPU. It creatively proposes a tensor expansion algorithm that can accurately adapt to the hardware features of the TPU, enabling the TPU to natively support tensor dimension expansion calculations, effectively filling the gap in TPU functionality, and significantly improving the computational parallelism and transmission efficiency in tensor expansion calculations.

[0199] Based on the above embodiments, the device may further include a memory alignment module, used for:

[0200] Before moving the data requiring dimensional expansion from the high-bandwidth memory of the TPU to the control memory or vector memory according to the tensor expansion requirements, memory alignment processing is performed on the lowest dimension data of the original tensor according to the data access granularity of the TPU.

[0201] Based on the above embodiments, the tensor expansion requirement may include at least one expansion dimension and an expansion dimension corresponding to each expansion dimension;

[0202] Accordingly, the extended processing module 330 may specifically include:

[0203] The basic information acquisition unit is used to sequentially acquire a target expansion dimension and the target expansion dimension corresponding to the target expansion dimension from the tensor expansion requirements in order from low dimension to high dimension.

[0204] The first extended processing unit is used to move the first data matching the target extended dimension from high-bandwidth memory to vector memory if the target extended dimension is the lowest dimension. Combined with the tensor broadcast mechanism on the TPU, the first extended data matching the target extended dimension is generated in vector memory through parallel processing of each core computing unit, and then the first extended data is moved back to high-bandwidth memory.

[0205] The second expansion processing unit is used to move the second data matching the target expansion dimension from the high-bandwidth memory to the control memory if the target expansion dimension is not the lowest dimension and the target expansion dimension is less than or equal to a preset dimension threshold. Combined with the tensor broadcast mechanism on the TPU, the second expansion data matching the target expansion dimension obtained by expanding the second data is directly stored back into the high-bandwidth memory through the parallel processing of each core computing unit.

[0206] The third extension processing unit is used to move the third data matching the target extension dimension from the high-bandwidth memory to the vector memory in blocks if the target extension dimension is not the lowest dimension and the target extension dimension is greater than the preset dimension threshold. Combined with the tensor broadcast mechanism on the TPU, the third extension data matching the target extension dimension is generated in blocks in the vector memory and then moved back to the high-bandwidth memory in blocks through the parallel processing of each core computing unit.

[0207] The repetitive execution unit is used to return the execution in order from low dimension to high dimension, sequentially obtaining a target expansion dimension and the corresponding target expansion dimension from the tensor expansion requirements, until all expansion dimensions are processed to obtain the target tensor data to be stored again in high-bandwidth memory.

[0208] Based on the above embodiments, a vector memory partitioning module may also be included, used for:

[0209] Before moving the data requiring dimensional expansion from the high-bandwidth memory of the TPU to the control memory or vector memory according to the tensor expansion requirements, the vector memory is divided into an input area and an output area. The input area is used to store the data in the original tensor, and the output area is used to store the data in the target tensor.

[0210] Based on the above embodiments, the first extended processing unit can be specifically used for:

[0211] The test checks whether the amount of data formed when expanding the first data to the first expanded data can be completely stored in the vector memory;

[0212] If the data can be stored completely, the first data is moved from the high-bandwidth memory to the input area of ​​the vector memory in one go, and the data expansion task that matches the target expansion dimension is split into two subtasks and distributed to the two core computing units.

[0213] Through two core computing units, using a tensor broadcast mechanism, TPU vector instructions are called in parallel to perform data copying operations on the assigned subtasks, forming the first extended data stored in the output area of ​​the vector memory, and the first extended data completely stored in the output area is moved back to the high-bandwidth memory in one go.

[0214] If it cannot be stored completely, determine the number of first blocks that match the first data, and move the first data to the input area of ​​the vector memory in batches according to the number of first blocks;

[0215] The data expansion task corresponding to the first data in each batch is split into two sub-tasks and distributed to two core computing units;

[0216] Through two core computing units, the TPU vector instructions are called in parallel using a tensor broadcast mechanism to perform data copying operations on the assigned subtasks, forming the first extended data of each batch stored in the output area of ​​the vector memory, and then the first extended data of each batch is moved back to the high-bandwidth memory in batches.

[0217] Based on the above embodiments, the second extended processing unit can be specifically used for:

[0218] The second data matching the target expansion dimension is moved from high-bandwidth memory to control memory in one go; the data expansion task matching the target expansion dimension is split into two sub-tasks and distributed to two core computing units;

[0219] Through two core computing units, a tensor broadcast mechanism is used to call TPU vector instructions in parallel to perform data copying operations on the assigned subtasks, and the expanded second data is then stored back into high-bandwidth memory.

[0220] Based on the above embodiments, the third extended processing unit can be specifically used for:

[0221] Determine the number of second blocks that match the third data, and move the third data to the input area of ​​the vector memory in batches according to the number of second blocks;

[0222] The data expansion task corresponding to the third data in each batch is split into two subtasks and distributed to two core computing units;

[0223] Through two core computing units, the TPU vector instructions are called in parallel using a tensor broadcast mechanism to perform data copying operations on the assigned subtasks, forming the third extended data of each batch stored in the output area of ​​the vector memory, and then the third extended data of each batch is moved back to the high-bandwidth memory in batches.

[0224] The TPU-based tensor expansion computing device provided in the embodiments of the present invention can execute the TPU-based tensor expansion computing method provided in any embodiment of the present invention, and has the corresponding functional modules and beneficial effects of the execution method.

[0225] The collection, storage, use, processing, transmission, provision, and disclosure of user personal information involved in the technical solution disclosed herein comply with the provisions of relevant laws and regulations and do not violate public order and good morals.

[0226] Figure 4 A schematic diagram of an electronic device 10, which can be used to implement embodiments of the present invention, is shown. The electronic device is intended to represent various forms of digital computers, such as laptop computers, desktop computers, workstations, personal digital assistants, servers, blade servers, mainframe computers, and other suitable computers. The electronic device can also represent various forms of mobile devices, such as personal digital processors, cellular phones, smartphones, wearable devices (e.g., helmets, glasses, watches, etc.), and other similar computing devices. The components shown herein, their connections and relationships, and their functions are merely illustrative and are not intended to limit the implementation of the invention described and / or claimed herein.

[0227] like Figure 4 As shown, the electronic device 10 includes at least one processor 11 and a memory, such as a read-only memory (ROM) 12 or a random access memory (RAM) 13, communicatively connected to the at least one processor 11. The memory stores computer programs executable by the at least one processor. The processor 11 can perform various appropriate actions and processes based on the computer program stored in the ROM 12 or loaded from storage unit 18 into the RAM 13. The RAM 13 can also store various programs and data required for the operation of the electronic device 10. The processor 11, ROM 12, and RAM 13 are interconnected via a bus 14. An input / output (I / O) interface 15 is also connected to the bus 14.

[0228] Multiple components in electronic device 10 are connected to I / O interface 15, including: input unit 16, such as keyboard, mouse, etc.; output unit 17, such as various types of displays, speakers, etc.; storage unit 18, such as disk, optical disk, etc.; and communication unit 19, such as network card, modem, wireless transceiver, etc. Communication unit 19 allows electronic device 10 to exchange information / data with other devices through computer networks such as the Internet and / or various telecommunications networks.

[0229] Processor 11 can be a variety of general-purpose and / or special-purpose processing components with processing and computing capabilities. Some examples of processor 11 include, but are not limited to, a central processing unit (CPU), a graphics processing unit (GPU), various special-purpose artificial intelligence (AI) computing chips, various processors running machine learning model algorithms, a digital signal processor (DSP), and any suitable processor, controller, microcontroller, etc. Processor 11 performs the various methods and processes described above, such as performing the TPU-based tensor expansion computing method as described in any embodiment of the present invention.

[0230] That is: to parse the tensor expansion request and obtain the first dimension information of the original tensor and the second dimension information of the target tensor;

[0231] Based on the first and second dimension information, determine the tensor expansion requirements;

[0232] Based on the tensor expansion requirements, the data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or vector memory. Combined with the tensor broadcasting mechanism on the TPU, the target tensor data is formed and stored back in the high-bandwidth memory through parallel processing of each core computing unit.

[0233] In some embodiments, the TPU-based tensor expansion computation method as described in any embodiment of the present invention can be implemented as a computer program tangibly contained in a computer-readable storage medium, such as storage unit 18. In some embodiments, part or all of the computer program can be loaded and / or installed on electronic device 10 via ROM 12 and / or communication unit 19. When the computer program is loaded into RAM 13 and executed by processor 11, one or more steps of the TPU-based tensor expansion computation method described above as described in any embodiment of the present invention can be performed. Alternatively, in other embodiments, processor 11 can be configured to perform the TPU-based tensor expansion computation method as described in any other suitable manner (e.g., by means of firmware).

[0234] Various embodiments of the systems and techniques described above herein can be implemented in digital electronic circuit systems, integrated circuit systems, field-programmable gate arrays (FPGAs), application-specific integrated circuits (ASICs), application-specific standard products (ASSPs), systems-on-a-chip (SoCs), payload-programmable logic devices (CPLDs), computer hardware, firmware, software, and / or combinations thereof. These various embodiments may include implementations in one or more computer programs that can be executed and / or interpreted on a programmable system including at least one programmable processor, which may be a dedicated or general-purpose programmable processor, capable of receiving data and instructions from a storage system, at least one input device, and at least one output device, and transmitting data and instructions to the storage system, the at least one input device, and the at least one output device.

[0235] Computer programs used to implement the methods of the present invention may be written in any combination of one or more programming languages. These computer programs may be provided to a processor of a general-purpose computer, a special-purpose computer, or other programmable data processing device, such that when executed by the processor, the computer programs cause the functions / operations specified in the flowcharts and / or block diagrams to be performed. The computer programs may be executed entirely on a machine, partially on a machine, or as a standalone software package, partially on a machine and partially on a remote machine, or entirely on a remote machine or server.

[0236] In the context of this invention, a computer-readable storage medium can be a tangible medium that may contain or store a computer program for use by or in conjunction with an instruction execution system, apparatus, or device. A computer-readable storage medium may include, but is not limited to, electronic, magnetic, optical, electromagnetic, infrared, or semiconductor systems, apparatus, or devices, or any suitable combination thereof. Alternatively, a computer-readable storage medium may be a machine-readable signal medium. More specific examples of machine-readable storage media include electrical connections based on one or more wires, portable computer disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), optical fibers, portable compact disk read-only memory (CD-ROM), optical storage devices, magnetic storage devices, or any suitable combination thereof.

[0237] To provide interaction with a user, the systems and techniques described herein can be implemented on an electronic device having: a display device (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) for displaying information to the user; and a keyboard and pointing device (e.g., a mouse or trackball) through which the user provides input to the electronic device. Other types of devices can also be used to provide interaction with the user; for example, feedback provided to the user can be any form of sensory feedback (e.g., visual feedback, auditory feedback, or tactile feedback); and input from the user can be received in any form (including sound input, voice input, or tactile input).

[0238] The systems and technologies described herein can be implemented in computing systems that include backend components (e.g., as data servers), or middleware components (e.g., application servers), or frontend components (e.g., user computers with graphical user interfaces or web browsers through which users can interact with implementations of the systems and technologies described herein), or any combination of such backend, middleware, or frontend components. The components of the system can be interconnected via digital data communication of any form or medium (e.g., communication networks). Examples of communication networks include local area networks (LANs), wide area networks (WANs), blockchain networks, and the Internet.

[0239] A computing system can include clients and servers. Clients and servers are generally located far apart and typically interact through communication networks. The client-server relationship is created by computer programs running on the respective computers and having a client-server relationship with each other. The server can be a cloud server, also known as a cloud computing server or cloud host, which is a hosting product within the cloud computing service system to address the shortcomings of traditional physical hosts and VPS services, such as high management difficulty and weak business scalability.

[0240] It should be understood that the various forms of processes shown above can be used, with steps reordered, added, or deleted. For example, the steps described in this invention can be executed in parallel, sequentially, or in different orders, as long as the desired result of the technical solution of this invention can be achieved, and this is not limited herein.

[0241] The specific embodiments described above do not constitute a limitation on the scope of protection of this invention. Those skilled in the art should understand that various modifications, combinations, sub-combinations, and substitutions can be made according to design requirements and other factors. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of this invention should be included within the scope of protection of this invention.

Claims

1. A tensor extension calculation method based on a tensor processing unit (TPU), characterized in that, The TPU includes two parallel core computing units and a hierarchical memory system including a high-bandwidth memory, a control memory and a vector memory. The tensor expansion request is parsed to obtain first dimension information of an original tensor and second dimension information of a target tensor; According to the first dimension information and the second dimension information, a tensor expansion requirement is determined; According to the tensor expansion requirement, data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or the vector memory, and the tensor broadcast mechanism on the TPU is combined to form target tensor data through parallel processing of each core computing unit and then the target tensor data is stored in the high-bandwidth memory.

2. The method of claim 1, wherein, Before moving the data in the original tensor that needs to be expanded in dimension from the high-bandwidth memory of the TPU to the control memory or the vector memory according to the tensor expansion requirement, the data in the lowest dimension of the original tensor is subjected to memory alignment processing according to the data access granularity of the TPU. The tensor expansion requirement includes at least one expansion dimension and an expansion dimension number corresponding to each expansion dimension; 3. The method according to claim 1 or 2, characterized in that, According to the tensor expansion requirement, data in the original tensor that needs to be expanded in dimension is moved from the high-bandwidth memory of the TPU to the control memory or the vector memory, and the tensor broadcast mechanism on the TPU is combined to form target tensor data through parallel processing of each core computing unit and then the target tensor data is stored in the high-bandwidth memory, including: In order from low dimension to high dimension, a target expansion dimension and a target expansion dimension number corresponding to the target expansion dimension are sequentially obtained from the tensor expansion requirement; If the target expansion dimension is the lowest dimension, first data matched with the target expansion dimension is moved from the high-bandwidth memory to the vector memory, and the tensor broadcast mechanism on the TPU is combined to generate first expansion data matched with the target expansion dimension number in the vector memory through parallel processing of each core computing unit, and then the first expansion data is moved back to the high-bandwidth memory; If the target expansion dimension is not the lowest dimension and the target expansion dimension number is less than or equal to a preset dimension threshold, second data matched with the target expansion dimension is moved from the high-bandwidth memory to the control memory, and the tensor broadcast mechanism on the TPU is combined to directly store second expansion data matched with the target expansion dimension number in the high-bandwidth memory through parallel processing of each core computing unit, the second expansion data being obtained by expanding the second data; If the target expansion dimension is not the lowest dimension and the target expansion dimension number is greater than the preset dimension threshold, third data matched with the target expansion dimension is moved from the high-bandwidth memory to the vector memory in blocks, and the tensor broadcast mechanism on the TPU is combined to generate third expansion data matched with the target expansion dimension number in the vector memory in blocks through parallel processing of each core computing unit, and then the third expansion data is moved back to the high-bandwidth memory in blocks; Return to execute the operation of sequentially obtaining a target expansion dimension and a target expansion dimension number corresponding to the target expansion dimension from the tensor expansion requirement in order from low dimension to high dimension until processing of all expansion dimensions is completed to obtain target tensor data stored in the high-bandwidth memory. ​ 4. The method of claim 3, wherein, Before the data required to be dimensionally expanded in the original tensor is carried from the high-bandwidth memory of the TPU to the control memory or vector memory according to the tensor expansion requirements, the following steps are further included: dividing the vector memory into an input area and an output area, wherein the input area is used to store data in the original tensor, and the output area is used to store data in the target tensor.

5. The method of claim 4, wherein, carrying first data matching the target expansion dimension from the high-bandwidth memory to the vector memory, and combining the tensor broadcast mechanism on the TPU, generating first expansion data matching the target expansion dimension in the vector memory through parallel processing of each core calculation unit, and then carrying the first expansion data back to the high-bandwidth memory, including: detecting whether the amount of data formed when the first data is expanded to the first expansion data can be completely stored in the vector memory; if it can be completely stored, carrying the first data from the high-bandwidth memory to the input area of the vector memory at one time, and splitting the data expansion task matching the target expansion dimension into two sub-tasks and distributing them to two core calculation units; through the two core calculation units, using the tensor broadcast mechanism and calling the TPU vector instruction in parallel to perform data replication operations on the allocated sub-tasks, and jointly forming the first expansion data stored in the output area of the vector memory, and then carrying the first expansion data completely stored in the output area back to the high-bandwidth memory at one time; if it cannot be completely stored, determining the first block quantity matching the first data, and carrying the first data to the input area of the vector memory in batches according to the first block quantity; splitting the data expansion task corresponding to each batch of first data into two sub-tasks and distributing them to two core calculation units; through the two core calculation units, using the tensor broadcast mechanism and calling the TPU vector instruction in parallel to perform data replication operations on the allocated sub-tasks, and jointly forming each batch of first expansion data stored in the output area of the vector memory, and then carrying each batch of first expansion data back to the high-bandwidth memory in batches.

6. The method of claim 3, wherein, carrying second data matching the target expansion dimension from the high-bandwidth memory to the control memory, and combining the tensor broadcast mechanism on the TPU, directly storing second expansion data matching the target expansion dimension obtained by expanding the second data back to the high-bandwidth memory through parallel processing of each core calculation unit, including: carrying the second data matching the target expansion dimension from the high-bandwidth memory to the control memory at one time; splitting the data expansion task matching the target expansion dimension into two sub-tasks and distributing them to two core calculation units; through the two core calculation units, using the tensor broadcast mechanism and calling the TPU vector instruction in parallel to perform data replication operations on the allocated sub-tasks, and storing the second expansion data obtained by expansion back to the high-bandwidth memory.

7. The method of claim 4, wherein, carrying third data matching the target expansion dimension from the high-bandwidth memory to the vector memory in batches, and combining the tensor broadcast mechanism on the TPU, generating third expansion data matching the target expansion dimension in the vector memory in batches through parallel processing of each core calculation unit, and then carrying the third expansion data back to the high-bandwidth memory in batches, including: determining a second number of blocks matching the third data, and carrying the third data in batches to an input area of the vector memory according to the second number of blocks; splitting a data expansion task corresponding to each batch of third data into two subtasks and distributing the two subtasks to two core computing units; using a tensor broadcasting mechanism to call a TPU vector instruction to perform a data copy operation on the allocated subtasks in parallel by the two core computing units, to jointly form each batch of third expansion data, and store the third expansion data to an output area of the vector memory, and carry back each batch of third expansion data to the high-bandwidth memory in batches.

8. An electronic device, comprising: The electronic device comprises: at least one processor; and a memory connected in communication with the at least one processor; wherein the memory stores a computer program executable by the at least one processor, and the computer program is executed by the at least one processor to enable the at least one processor to execute the TPU-based tensor expansion calculation method of any one of claims 1-7.

9. A computer-readable storage medium, characterized in that, The computer readable storage medium stores computer instructions for causing the processor to implement the TPU-based tensor expansion calculation method of any one of claims 1-7 when executed.

10. A computer program product, characterised in that, The computer program product comprises a computer program which, when executed by a processor, implements the TPU-based tensor expansion calculation method according to any one of claims 1-7. The computer program product comprises a computer program which, when executed by a processor, implements the TPU-based tensor expansion calculation method according to any one of claims 1-7.

Images