Tensor dimension reorganization method and device for tensor processing unit, chip
By employing a high-frequency special judgment mode and a hierarchical storage structure in the TPU, the problem of bandwidth resource waste and low hardware utilization in the TPU reshape operation is solved, achieving efficient tensor dimension reorganization and improving computational efficiency and adaptability.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- ZHONGHAO XINYING (HANGZHOU) TECH CO LTD
- Filing Date
- 2026-01-22
- Publication Date
- 2026-05-12
AI Technical Summary
Existing Tensor Processing Units (TPUs) suffer from wasted bandwidth resources and low hardware utilization when performing reshape operations, especially in combinations of operators involving multiple reshapes, transposes, or concatenations. This leads to intermediate results being written back and repeatedly rearranged, undermining the high-throughput pipeline originally designed for matrix computations.
By implementing a tensor dimension reorganization method in a TPU, employing a high-frequency special judgment mode and a hierarchical storage structure, and utilizing vector parallel computing units and hardware acceleration execution, tensor dimension reorganization is completed within the on-chip storage range, avoiding redundant data movement and complex index calculations. A multi-processing unit collaborative working mechanism is used for pipelined processing.
It improves the execution efficiency of reshape operations, reduces on-chip memory resource consumption, avoids bandwidth resource waste, improves hardware utilization, and adapts to complex model structures, thereby enhancing computational efficiency and adaptability.
Smart Images

Figure CN121542216B_ABST
Abstract
Description
Technical Field
[0001] This application relates to the field of tensor processing, and more specifically, to a method, apparatus, and chip for tensor dimension reorganization in a tensor processing unit. Background Technology
[0002] With the development of artificial intelligence and deep learning technologies, the reshape operator is frequently used in computation graphs to adjust the shape of tensors to adapt to the input and output requirements of different operators. This operation does not change the total amount of data, but only restructures the dimensional structure, and is a fundamental component of deep learning frameworks.
[0003] However, the Tensor Processing Unit (TPU) currently does not natively support reshape. Adopting the CPU or GPU implementation that rearranges data element-wise using scalar cores and writes it back would lead to the following technical problems: First, the TPU's on-chip memory employs a multi-level buffer and multi-bank architecture, optimized for high-throughput matrix multiplication. Direct porting would result in numerous non-sequential and cross-bank random accesses, increasing the number of data moves and access conflicts between on-chip and high-bandwidth memory, severely wasting bandwidth resources. Second, existing CPU or GPU implementations often treat reshape as an independent operator, triggering a complete data rearrangement process with each shape transformation. In combinations of operators involving multiple reshapes, transposes, or concatenations, multiple intermediate result write-backs and repeated rearrangements are likely to occur, leading to excessive intermediate buffer space usage. This undermines the TPU's original high-throughput pipeline designed for matrix computation and reduces the hardware utilization of the computing unit.
[0004] Therefore, there is an urgent need to propose a technical solution to implement tensor reshape operations in TPU and solve at least one of the above-mentioned technical problems. Summary of the Invention
[0005] In this context, embodiments of this application aim to provide a method, apparatus, and chip for tensor dimension reshaping in a tensor processing unit, which can implement tensor reshape operations in a TPU, avoid bandwidth resource waste, and improve hardware utilization.
[0006] In a first aspect of the embodiments of this application, a method for tensor dimension reorganization for a tensor processing unit is provided, the tensor processing unit having a hierarchical storage structure and a vector parallel computing unit, the method comprising:
[0007] For the input tensor to be processed in the tensor processing unit, it is verified whether the total number of elements of the input tensor and the target output tensor are equal to meet the basic constraints of tensor dimension reorganization.
[0008] Based on the shape parameters and storage layout information of the input tensor and the target output tensor, as well as the hardware architecture characteristics of the tensor processing unit, a mode determination is performed on the current dimension recombination operation to determine whether the current dimension recombination operation matches a preset high-frequency special judgment mode.
[0009] When the current dimension recombination operation matches the high-frequency special judgment mode, the hardware acceleration execution method corresponding to the high-frequency special judgment mode is adopted to complete the tensor dimension recombination within the on-chip storage range and obtain the target output tensor.
[0010] When the current dimension reorganization operation does not match the high-frequency special judgment mode, a general reorganization execution method oriented towards hierarchical storage structure is adopted. Tensor dimension reorganization is completed through parallel computing and collaborative processing to obtain the target output tensor.
[0011] In a second aspect of this application, a tensor dimension reorganization apparatus for a tensor processing unit is provided. The apparatus is applied to the tensor processing unit, which has a hierarchical storage structure and a vector-parallel computation unit. The apparatus includes the following units:
[0012] The input unit is used to verify whether the total number of elements of the input tensor and the target output tensor are equal to the input tensor to be processed, so as to meet the basic constraints of tensor dimension reorganization.
[0013] The judgment unit is used to perform mode determination on the current dimension recombination operation based on the shape parameters and storage layout information of the input tensor and the target output tensor, as well as the hardware architecture characteristics of the tensor processing unit, so as to determine whether the current dimension recombination operation matches the preset high-frequency special judgment mode.
[0014] The processing unit is used to, when the current dimension reorganization operation matches the high-frequency special judgment mode, adopt a hardware-accelerated execution method corresponding to the high-frequency special judgment mode to complete the tensor dimension reorganization within the on-chip storage range and obtain the target output tensor; when the current dimension reorganization operation does not match the high-frequency special judgment mode, adopt a general reorganization execution method oriented towards hierarchical storage structure, and complete the tensor dimension reorganization through parallel computing and collaborative processing to obtain the target output tensor.
[0015] In a third aspect of the embodiments of this application, a terminal device is provided, the terminal device comprising: at least one processor, a memory, and an input / output unit; wherein the memory is used to store a computer program, and the processor is used to invoke the computer program stored in the memory to execute the tensor dimension reorganization method for a tensor processing unit as described in the first aspect.
[0016] In a fourth aspect of the embodiments of this application, a computer-readable storage medium is provided, which includes instructions that, when executed on a computer, cause the computer to perform the tensor dimension reorganization method for a tensor processing unit as described in any one aspect.
[0017] In a fifth aspect of the embodiments of this application, a computer program product is provided, including a computer program that, when executed by a processor, implements the tensor dimension reorganization method for a tensor processing unit as described in the first aspect.
[0018] In a sixth aspect of the embodiments of this application, a chip is provided that includes a computer program, which, when executed by a processor, implements the tensor dimension reorganization method for a tensor processing unit as described in the first aspect.
[0019] This application discloses a method, apparatus, and chip for tensor dimension reorganization in a tensor processing unit. The technical solution of this application is applied to a tensor processing unit, which has a hierarchical storage structure and a vector parallel computing unit. First, for the input tensor to be processed, the total number of elements in the input tensor and the target output tensor are checked to ensure they meet the basic constraints of tensor dimension reorganization. Then, based on the shape parameters and storage layout information of the input and target output tensors, as well as the hardware architecture characteristics of the tensor processing unit, a mode determination is performed on the current dimension reorganization operation to determine whether it matches a preset high-frequency special judgment mode. When the current dimension reorganization operation matches the high-frequency special judgment mode, a hardware-accelerated execution method corresponding to the high-frequency special judgment mode is used to complete the tensor dimension reorganization within the on-chip storage range. When the current dimension reorganization operation does not match the high-frequency special judgment mode, a general reorganization execution method oriented towards the hierarchical storage structure is used to complete the tensor dimension reorganization through parallel computing and collaborative processing.
[0020] This application's technical solution utilizes special pattern recognition and hardware acceleration mechanisms to achieve tensor dimension reorganization operations for tensor processing units, thereby improving the reshape execution efficiency of tensor processing units. This solution separates dimension transformation patterns (such as single-dimensional transpose and conversion between dimensions at specific ratios) from the general data reorganization process, directly completing the transformation through dedicated hardware instructions, avoiding redundant data movement and complex index calculations. Simultaneously, through a multi-processing unit collaborative working mechanism, data loading, reorganization, and storage operations are pipelined, effectively shortening storage access latency. This not only significantly improves the execution efficiency of reshape operations but also reduces on-chip memory resource consumption, avoids bandwidth waste, improves hardware utilization, and provides efficient support for dynamic tensor operations in various machine learning models, enabling tensor processing units to exhibit stronger adaptability and computational efficiency when facing complex model structures. Attached Figure Description
[0021] Figure 1 This application illustrates a flowchart of a tensor dimension reorganization method for a tensor processing unit.
[0022] Figure 2 This is a schematic diagram of a tensor dimension reorganization device for a tensor processing unit shown in this application. Detailed Implementation
[0023] The following is for reference. Figure 1 , Figure 1 This is a flowchart illustrating a method, apparatus, and chip for tensor dimension reorganization in a tensor processing unit, provided as an embodiment of this application.
[0024] To address at least one technical problem existing in the related technologies, embodiments of this application provide a method, apparatus, and chip for tensor dimension reorganization in a tensor processing unit.
[0025] Specifically, this application's embodiments improve the execution efficiency of common reshape operations through accurate identification and dedicated optimization of high-frequency special judgment patterns. In large model inference scenarios, the general reshape implementation in related technologies requires complex index calculations and multiple data moves, while this application's embodiments can directly identify such high-frequency patterns and utilize the hardware transpose capability of the TPU vector register to complete data rearrangement within the on-chip storage range, avoiding redundant data transmission between different levels of storage.
[0026] For the TPU's unique hierarchical storage architecture (HBM / VMEM / CMEM), this application proposes an adaptive memory partitioning strategy and data alignment mechanism to maximize the bandwidth characteristics of each storage level during reshape operations. By dynamically calculating the VMEM block size, it ensures precise matching between the processing granularity of data blocks and vector computation units, reducing gaps and waste in memory access. In the general reshape process, a dual-processing unit collaborative mechanism is adopted. One unit is responsible for loading data from HBM to VMEM and caching it in CMEM, while the other unit simultaneously reads data from CMEM, reassembles it, and writes it back to HBM. This pipelined design effectively masks storage access latency, improves VMEM bandwidth utilization, and increases CMEM cache hit rate, fundamentally solving the memory wall problem in reshape operations.
[0027] This application utilizes the TPU's 8×128 vector parallel computing capability to refactor the traditional serial reshape operation into vector-level parallel processing. Through vector concatenation (cat128) and vector splitting (get128) instructions, 1024 data elements can be rearranged in a single clock cycle. In specific dimension conversion scenarios (such as the conversion between 320 and 40), the vector rotation instruction (m_rotate) is applied to directly adjust the element order at the register level, avoiding the complex index calculations and conditional branches of traditional methods. This vector-level operation not only significantly improves computational efficiency but also reduces the complexity of control logic, enhancing the utilization of the TPU's vector computing resources and fully leveraging the hardware's parallel potential while ensuring data consistency.
[0028] For reshape operations with unchanged physical layout, this application implements an intelligent same-layout recognition mechanism that can accurately detect the consistency of the distribution of input and output tensors in physical memory. When such a scenario is identified, all computational units are bypassed, and the operation is completed through high-bandwidth memory DMA direct transfer, completely avoiding unnecessary data parsing and reorganization computations. This zero-computation optimization strategy is valuable in common deep learning operations such as feature concatenation and dimensionality compression, reducing invalid computational overhead and lowering system power consumption, providing important support for energy-sensitive edge computing scenarios.
[0029] The pattern determination mechanism in this application can dynamically select the optimal execution path based on the shape parameters, storage layout, and hardware characteristics of the input and target tensors. Whether it's an arbitrary dimensional conversion from 1D to 5D or the processing requirements of different data types, it can automatically match the most suitable processing strategy. This adaptive capability is particularly suitable for dynamic computation graph execution environments, such as variable batch sizes or dynamic shape tensors in large model inference, providing broad scenario compatibility while maintaining high performance.
[0030] From a system architecture perspective, this embodiment deeply integrates the reshape operation into the TPU hardware architecture, reducing communication overhead and synchronization wait times between the host CPU and the TPU device. Traditional reshape implementations typically require CPU involvement in index calculation and task scheduling, while this embodiment completely offloads the core logic to the TPU device, enabling autonomous device execution, reducing host load, and improving overall system throughput. In end-to-end deep learning model execution, this optimization reduces host-device communication latency and achieves a more balanced allocation of computing resources, providing a solid foundation for building larger-scale and more complex AI models. Especially in generative AI and large-model training scenarios, frequent dimensionality transformation operations no longer become a performance bottleneck, improving overall training efficiency and accelerating the iteration and deployment of AI models.
[0031] Figure 1 The flowchart of a tensor dimension reorganization method for a tensor processing unit, as shown in one embodiment of this application, includes:
[0032] Step S101: For the input tensor to be processed in the tensor processing unit, verify whether the total number of elements of the input tensor and the target output tensor are equal, so as to satisfy the basic constraint of tensor dimension reorganization.
[0033] This application applies to a tensor processing unit. The tensor processing unit has a hierarchical storage structure and a vector-parallel computation unit. In the dimension reshaping operation of the tensor processing unit, the input tensor represents the original multidimensional data structure to be processed, which has specific dimensional shapes, data types, and memory layout characteristics. The input tensor is typically stored in high-bandwidth memory (HBM) as the source data carrier for the reshape operation. The target output tensor is the expected result data structure of the reshape operation. The target output tensor retains all data elements of the input tensor but is organized according to a new dimensional configuration to meet the specific needs of subsequent computational tasks.
[0034] In this embodiment, the tensor processing unit utilizes its hierarchical storage architecture (including high-bandwidth memory, on-chip memory, and cache) and vector parallel computing capabilities to efficiently process these two types of tensors. The most critical constraint is that the total number of elements in the input tensor and the target output tensor must be strictly equal. This fundamental constraint ensures the mathematical validity and data integrity of the reshape operation. When this constraint is not met, the dimension reorganization operation cannot be executed, and the system returns an error. When the constraint is met, the tensor processing unit selects a high-frequency special judgment mode or a general reorganization execution path based on the specific dimensional characteristics of the input and output tensors. Through optimized data flow and computational strategies, it achieves efficient dimensional structure transformation while maintaining the data content. This mechanism is particularly suitable for frequent dimensional transformation operations in large model inference and training scenarios, such as single-dimensional transpose and conversion between 320 and 40 dimensions, significantly improving tensor processing efficiency and reducing computational latency.
[0035] Understandably, the three-level hierarchical storage architecture adopted by the tensor processing unit provides an efficient hardware foundation for tensor dimension reorganization operations by having storage media with different performance characteristics work together.
[0036] In the hierarchical storage architecture of the Tensor Processing Unit (TPU), High Bandwidth Memory (HBM), as the largest storage layer, is located outside the TPU chip and is responsible for storing large-scale input and output tensor data. Its main function is to provide sufficient storage space and high total bandwidth to meet the massive data demands of large model processing. On-chip Memory (VMEM), as an intermediate layer, is directly integrated inside the TPU chip. It has lower access latency than HBM and is mainly used for temporary storage of data blocks being processed. While its capacity is limited, its access speed is faster, acting as a data transfer station during dimension reorganization. Especially for high-frequency special cases, VMEM can directly complete dimension transformation without further access to higher-level caches. Cache (CMEM), located at the top of the storage hierarchy, has the fastest access speed but the most limited capacity. It is mainly used to cache frequently accessed small-scale data, such as intermediate results in vector computation.
[0037] In this embodiment, these three levels of storage are efficiently interconnected through a Direct Memory Access (DMA) engine, and the data flow path can be dynamically optimized according to the processing scenario: For general reshape scenarios, a complete path of HBM→VMEM→CMEM→VMEM→HBM is used, fully utilizing the characteristics of each level of storage. For special cases, the path is intelligently tailored; for example, transpose reshape directly completes the dimension transformation in VMEM, and same-layout reshape even bypasses on-chip storage to achieve direct transmission from HBM to HBM. This hierarchical design works deeply with the TPU's 8×128 vector computing units, ensuring efficient data flow between storage levels through strict memory alignment strategies (such as 128-element granularity alignment) and dedicated data transport instructions (such as dlc_dma). This greatly improves the execution efficiency of reshape operations in large model inference scenarios, especially when processing qwen. 72b It exhibits performance advantages when modeling large-scale models.
[0038] Furthermore, in the hierarchical storage architecture of the Tensor Processing Unit (TPU), the storage of input tensors and target output tensors follows the principle of data locality and computational efficiency optimization strategies, and its storage distribution has a clear hierarchical characteristic.
[0039] The initial storage locations of the input and target output tensors are typically initially located in the HBM, a large off-chip storage area of the TPU with high storage capacity (usually reaching GB levels) but relatively high access latency. The data layout in the HBM follows a specific memory format, including dimension order, padding regions, and alignment requirements, which directly affect subsequent processing efficiency. For example, when the input tensor is of shape [256, 1], its physical storage in the HBM may include additional padding to meet the hardware requirement of 128-element alignment.
[0040] During the dimension reorganization operation, this application further provides a method for data to dynamically flow between storage levels. For example, in a general processing scenario, the data flow path is HBM→VMEM→CMEM→VMEM→HBM. The input tensor is first loaded from HBM to on-chip memory (VMEM) via DMA, and then the processing unit with device_id=0 further caches the data block to the cache (CMEM). After the reorganization calculation is completed, the processing unit with device_id=1 reads the result from CMEM, temporarily stores it in VMEM, and finally writes it back to HBM to form the target output tensor. In the transpose-based special case scenario, a simplified path of HBM→VMEM→HBM is adopted, and the dimension transformation is directly completed in VMEM, avoiding the interaction overhead with CMEM. For example, the operation of transposing [C, 1] to [1, C] is completed in VMEM through a register-level transpose instruction. In the same layout direct copy scenario, VMEM and CMEM are completely bypassed, and data copying with zero computational overhead is achieved through direct DMA transfer from HBM to HBM. In a single-dimensional unfolding scenario, it may be possible to modify only the tensor metadata (dimensional structure, stride information) while keeping the physical data position unchanged, thus achieving a "view transformation" with zero data movement.
[0041] TPUs improve data processing efficiency by controlling the storage hierarchy. For example, when allocating buffers in VMEM, a 16KB boundary alignment strategy is used to divide the available space into input and output block buffers, ensuring DMA transfer efficiency. In CMEM, data is organized at a 128-element vector granularity, adapting to the width of the TPU's vector processing units, enabling efficient execution of vector instructions such as cat128 / get128. The shape description information of tensors (number of dimensions, size, stride, storage offset, etc.) is usually cached in dedicated registers to avoid frequent access to HBM. In parallel processing scenarios, input tensors are divided into multiple data blocks, each processed by different processing units (e.g., device_id=0 and device_id=1), with each unit having an independent data slice in its own VMEM region.
[0042] Furthermore, for high-frequency special cases, the TPU employs special storage strategies. For example, shared storage views mean that in single-dimensional expansion optimization, the input tensor and the target output tensor may share the same physical storage area, achieving logical dimensional changes only by modifying metadata, avoiding actual data movement. Another example is a prefetch caching mechanism, which handles specific dimensional conversions (such as 320). 40) The system prefetches adjacent data blocks into the VMEM to reduce the number of HBM accesses by leveraging data locality. For example, boundary mask control uses a vector mask mechanism to identify valid data regions when processing unaligned data blocks, avoiding the storage access overhead of invalid data.
[0043] This hierarchical, scenario-adaptive storage mechanism enables the TPU to maximize the execution efficiency of dimension reorganization operations while ensuring data integrity, making it particularly suitable for various reshape scenarios that frequently occur in large model inference.
[0044] As an optional embodiment, in step S101, verifying whether the total number of elements in the input tensor and the target output tensor are equal to satisfy the basic constraint of tensor dimension reorganization includes: obtaining the number of dimensions and the size of each dimension from the tensor description information corresponding to the input tensor and the target output tensor respectively; calculating the total number of elements in the input tensor based on the size of each dimension of the input tensor; calculating the total number of elements in the target output tensor based on the size of each dimension of the target output tensor; if the total number of elements in the input tensor is consistent with the total number of elements in the target output tensor, then performing the tensor dimension reorganization operation.
[0045] In the above embodiment, in step S101, firstly, the number of dimensions and the size of each dimension are obtained from the tensor description information corresponding to the input tensor and the target output tensor, respectively. The tensor description information is stored in a dedicated register of the tensor processing unit and includes at least the number of dimensions, the size of each dimension, the dimension stride, the storage offset, and dimension padding information. Subsequently, based on the dimensions of each dimension of the input tensor, the total number of elements in the input tensor, in_numel, is calculated using a hardware multiplier. The calculation formula could be, for example, in_numel = shape[0] × shape[1] × shape[2] × shape[3] × shape[4], where shape[i] represents the size of the i-th dimension. This calculation process supports any combination of dimensions from 1D to 5D tensors and uses 64-bit integer operations to prevent overflow. Similarly, based on the dimensions of each dimension of the target output tensor, the total number of elements out_numel of the target output tensor is calculated. During the calculation process, if any tensor is detected to have a dimension padding region, the elements in the padding region are automatically excluded based on the padding identifier in the tensor description information, and only the valid data region is calculated. Furthermore, the verification operation also includes checking the data type consistency between the input tensor and the target output tensor. When the data types are different, the calculation basis for the total number of elements is adjusted according to a preset type conversion rule. If the total number of elements in the input tensor is consistent with the total number of elements in the target output tensor, a tensor dimension reorganization operation is performed. If they are inconsistent, the processing flow is terminated, a dimension mismatch error code is returned to the caller, and the allocated on-chip memory resources are released.
[0046] In practical applications, optionally, the total number of elements verification operation is executed by the scalar processing core of the tensor processing unit during the task scheduling phase, avoiding the occupation of vector computing resources and improving overall processing efficiency.
[0047] Optionally, after step S101, the hardware parameter configuration phase of the tensor processing unit includes the allocation and alignment optimization of on-chip memory resources. Specifically, the base addresses of the input and output tensors in high-bandwidth memory are first initialized, and then a dedicated data processing area is configured in the on-chip memory (VMEM). This configuration process is implemented through the following steps: The starting address of the on-chip memory provided by the tensor processing unit, info->vmem_addr, is obtained, and the corresponding tensor object vmem is created. Then, the available on-chip memory length vmemlen is calculated, which is the smaller value between info->vmem_size / 4 and 3MB (1024×1024×3), and the memory boundary is aligned to 128 bytes by bitwise operations (-128) to match the 128-element processing granularity of the TPU vector processing unit. At the same time, the starting address of the cache (CMEM), info->cmem_addr, is obtained, and the cmem object is initialized to provide high-speed temporary storage space for subsequent vector-level data reassembly.
[0048] Regarding data alignment configuration, this application can further employ template metaprogramming technology, defining dedicated alignment rules for data types with different bit widths through the BitWidthInfo template class. Taking a 32-bit data type as an example, the system configures the AddrShift parameter to 5 (representing 2). 5 =32, corresponding to the address offset of a 4-byte element), and implements the Align function to align any length l upwards to the 128-byte boundary (calculated by (l+127)&0xffffff80). This alignment strategy is particularly well-suited to the TPU's vector instruction set, such as the vector loading instruction vld and the vector concatenation instruction cat128, which can improve memory access efficiency.
[0049] Furthermore, the memory partitioning strategy considers the needs of multi-core collaborative processing, dividing the VMEM space into multiple independent regions based on device_id to avoid data contention between cores. Simultaneously, the alignment mechanism also includes boundary safety checks to prevent out-of-bounds memory access. In high-frequency special judgment scenarios, this refined memory configuration can improve data transmission bandwidth, especially suitable for large models (such as qwen). 72b The dimension transformation operation appears frequently in the code. It is worth noting that the VMEM length is limited to, for example, 3MB. This is an optimized selection based on the physical resources of the TPU chip, ensuring sufficient memory space is reserved for the execution of other cooperative operators, reflecting the global optimization idea of resource allocation in this application.
[0050] Step S102: Based on the shape parameters and storage layout information of the input tensor and the target output tensor, as well as the hardware architecture characteristics of the tensor processing unit, a mode determination is performed on the current dimension recombination operation to determine whether the current dimension recombination operation matches a preset high-frequency special judgment mode.
[0051] The shape parameters of the input and target output tensors include, but are not limited to, the number of dimensions, the size of each dimension, and the order of the dimensions. Storage layout information includes, but is not limited to, stride arrays, storage offsets, padding markers, and memory continuity markers. In practical applications, mathematical features of dimensional transformations are extracted, such as the dimension swap ratio (320 / 40=8), single-dimensional structure identifiers (dimension size of 1), and the complexity index of the linear index mapping function. For example, when the input tensor shape is detected as [C, 1] and the target output is [1, C], it is further verified whether C is an integer multiple of 128 to determine if the hardware alignment requirements are met.
[0052] The physical architecture characteristics of the tensor processing unit include, but are not limited to: the 128-element processing granularity of the vector register, the 3MB effective capacity limit of the on-chip memory (VMEM), the access latency characteristics of the cache (CMEM), the 128-byte transfer alignment requirement of the DMA engine, and the cooperative mechanism of the multi-core processing unit. These hardware parameters constitute the feasibility boundary of pattern matching, such as the limitation that the single data block processed when transforming from a high dimension to a low dimension does not exceed the VMEM capacity.
[0053] As an optional embodiment, in step S102, when performing pattern determination on the current dimension recombination operation to determine whether the current dimension recombination operation matches a preset high-frequency special judgment pattern, at least one of the following methods is included:
[0054] Method 1: Obtain the dimensions and dimensional order of the input tensor and the target output tensor, and determine whether the input tensor and the target output tensor satisfy a preset one-to-one correspondence or interchange relationship between dimensions.
[0055] In Method 1, the main approach involves parsing the dimension descriptors of the input and target output tensors to obtain the arrays of dimensions and their order, calculating the dimension mapping matrix, and determining if a predefined dimension correspondence exists. For example, if the input tensor dimension is [C, 1] and the target output tensor dimension is [1, C], and C is a multiple of 128, it is determined to meet the transpose recombination pattern. Alternatively, in another example, if an 8:1 ratio between the input dimension 320 and the output dimension 40 is detected, and the conditions in->storage_offset=0 and out->storage_offset=0 are met, 320 is triggered. 40 specific dimension conversion modes.
[0056] Method 2: Determine whether there is a single-dimensional structure with a dimension size of a preset threshold in the input tensor and / or the target output tensor, and determine whether the shape constraint condition of the high-frequency special judgment mode is satisfied based on the position of the single-dimensional structure in the tensor.
[0057] Method 2 primarily targets single-dimensional structures for specific detection. It iterates through the dimensions of the input and output tensors, identifies single dimensions with a value of 1, and records their position index in the dimension sequence. If the first dimension of the input tensor has a size of 1 and the second dimension does not exceed the cache capacity threshold (e.g., 1024×1024×3 bytes), and the data type is dlc_int32 or dlc_fp32, the single-dimensional expansion optimization path is activated, and the physical layout is reconstructed in on-chip memory using the transpose_in function.
[0058] Optionally, during the register prefetching phase, when the task scheduler assigns the reshape operation to the TPU, the scalar processing core automatically loads the TensorDesc structure from HBM into a dedicated configuration register set. This process is executed in parallel with DMA data transfer, hiding memory access latency. For example, all dimension information can be loaded at once using the ld.desc r0, [in_tensor_ptr-64] instruction. During the on-chip cache verification phase, it is also possible to check whether the dimension size meets the pre-screening conditions of the high-frequency special judgment mode, such as verifying REG_SHAPE[1]==1 (detecting the [C,1] structure) or REG_SHAPE[0]% 128 == 0 (128 alignment check). These scalar operations are completed within 2 clock cycles without accessing on-chip memory (VMEM). During the vector register mapping phase, for multi-dimensional judgment scenarios (such as the conversion between 320 and 40), the dimension size is broadcast to the vector register set, and mode matching is accelerated by single instruction multiple data (SIMD) comparison. For example, the command vcmp.eq v0, REG_SHAPE[0], #320 can simultaneously validate multiple dimensions of conditions.
[0059] Method 3: Obtain the storage layout information of the input tensor and the target output tensor. The storage layout information includes at least the storage offset, padding method and continuity identifier. Based on the storage layout information, determine whether the input tensor and the target output tensor are consistent in physical storage layout.
[0060] In Method 3, the storage layout information is analyzed in depth to extract storage offsets (inoff.lo, outoff.lo), dimension padding identifiers (in->dim0_padding), continuity markers, and memory step size arrays from tensor metadata. When all offsets are zero, there is no dimension padding, and the total number of elements is equal, it is determined to be a direct copy scenario with the same layout. The calculation and rearrangement steps will be skipped, and direct high-bandwidth memory to high-bandwidth memory transfer will be achieved directly through the dma_h2h instruction.
[0061] Method 4: Determine whether the data types of the input tensor and the target output tensor are consistent, and whether they meet the data width and alignment conditions required by the vector instruction of the tensor processing unit.
[0062] Method 5: Determine whether it belongs to the set of data types that can be directly recombined and supported by the tensor processing unit.
[0063] Methods four and five together achieve data type adaptation. First, it verifies whether the data type identifiers of the input and output tensors are consistent. Then, it checks whether the type belongs to the predefined set {dlc_fp32, dlc_int32, dlc_fp16} supported by the TPU vector instruction set, and verifies whether the memory address meets the alignment requirements corresponding to the type (e.g., 32-bit data requires 128-byte alignment). This process is implemented through the AddrShift parameter and Align function of the BitWidthInfo template class.
[0064] For example, data type determination begins by extracting the type identifier from the metadata region of the input and output tensors. In the TPU architecture, each tensor descriptor header contains an 8-bit type identifier field (dtype_id), which maps to predefined enumerated values such as dlc_fp32=0x00, dlc_int32=0x01, dlc_fp16=0x02, dlc_int8=0x03, etc. For instance, the determination process is implemented by a dedicated comparison circuit: First, the dtype_ids of the input and output tensors are obtained using the scalar loading instructions ld.scalar r0, [in_tensor_ptr-4] and ld.scalar r1, [out_tensor_ptr-4], respectively. Then, the hardware comparison instruction cmp.eq r2, r0, r1 is executed. If the result is true (r2=1), the data types are consistent, and alignment condition checks are performed. Otherwise, the special pattern matching is immediately terminated, and the process falls back to the general processing path. This process can be completed within one clock cycle, further reducing latency.
[0065] Exemplarily, when the data types are the same, further verify whether the type belongs to the direct recombination set natively supported by the TPU vector instruction set. This verification is achieved by looking up the type support vector register (TYPE_SUPPORT_VEC), which is a 32-bit bitmap, and each bit corresponds to the support status of a data type. For example, the bitmap value 0x00000007 (binary...00000111) indicates support for the first three types (dlcfp32 / dlc_int32 / dlc_fp16). The verification process based on this is as follows: calculate the mask mask = 1<<dtype_id, perform a bit test bt r3, TYPE_SUPPORT_VEC, dtype_id, and if r3 = 1, the type is supported. For dlc_fp32 (dtype_id = 0), the mask is 0x00000001, and the bit test result is true; for dlc_int8 (dtype_id = 3), the mask is 0x00000008, and if the corresponding bit in the bitmap is 0, the test fails. This mechanism allows the firmware to dynamically update the supported type set without modifying the hardware logic.
[0066] Thus, among the above five methods, when at least one determination result meets the trigger condition corresponding to the high-frequency special determination mode, the current dimension recombination operation is determined to match the high-frequency special determination mode.
[0067] Exemplarily, manage each special determination mode by setting a priority queue. When multiple determination conditions are met simultaneously, select the highest-priority mode in the order of transposed recombination (priority 0), same-layout direct copy (priority 1), vector width conversion (priority 2), and single-dimension expansion (priority 3). If any method returns a matching result, immediately terminate the subsequent determination, classify the current operation as the corresponding high-frequency special determination mode, and record the matching mode identifier for use in the subsequent execution phase.
[0068] Step S103, when the current dimension recombination operation matches the high-frequency special determination mode, adopt the hardware acceleration execution method corresponding to the high-frequency special determination mode to complete the tensor dimension recombination within the on-chip storage range.
[0069] As an optional embodiment, in step S103, the dimensional correspondence between the input tensor and the target output tensor is obtained, and it is determined that the current dimension reorganization operation belongs to the transpose reorganization scenario. The triggering conditions for the transpose reorganization scenario include: the input tensor and the target output tensor are two-dimensional tensors, the first and second dimensions of the input tensor and the target output tensor are interchanged, and at least one dimension is a single-dimensional structure with a preset threshold size. Then, if it is determined that the preset alignment granularity and data type constraints of the vector parallel computing unit of the tensor processing unit are satisfied, storage space is allocated in the on-chip memory for the input block buffer and the output block buffer. The storage space allocation is performed entirely in the on-chip memory and does not involve data caching operations in the cache. Next, the input tensor is divided into blocks according to the capacity of the on-chip memory, and each block is moved from the high-bandwidth memory to the input block buffer through direct memory access, avoiding loading data into the cache. Then, the vector-level transpose and reordering instruction or register-level transpose and rearrangement mechanism of the vector parallel computing unit is invoked to arrange each block in the input block buffer according to the dimension of the target output tensor and write it into the output block buffer. The dimension transformation is completed directly in on-chip memory, bypassing the intermediate processing of the cache. Finally, the reordering result in the output block buffer is written back to the target output tensor in high-bandwidth memory through direct memory access.
[0070] Specifically, when step S103 detects that the current dimension recombination operation matches a transpose recombination scenario, a highly optimized hardware acceleration process is executed. The triggering conditions for the transpose recombination scenario can be set as follows: both the input tensor and the target output tensor are two-dimensional structures, and they satisfy the dimension interchange relationship (in->dim0 == out->dim1 and in->dim1 == out->dim0), where at least one dimension is a single-dimensional structure (in->dim1 == 1), the main dimension size is an integer multiple of the vector processing unit granularity ((in->dim0&127) == 0), and the data type is 32-bit floating-point (in->dtype == dlc_fp32&&out->dtype == dlc_fp32). When the above triggering conditions are met, memory allocation is performed in on-chip memory (VMEM). First, the available VMEM capacity is calculated. System area is reserved using `_VMEMsize = MAX_VMEMsize - (MAX_VMEMsize / 128)`, and then 16KB boundary alignment is ensured using `_VMEMsize = _VMEMsize & -16384`. Subsequently, the VMEM space is divided into two contiguous regions. `CxxTensor` (C++ tensor object constructor) `vmem_output(info->vmem_addr)` defines the starting position of the output buffer, and `CxxTensor vmem_input(vmem_output + (_VMEMsize / 32))` defines the input buffer, with a ratio of 31:1 to optimize data flow. During the data transfer phase, block transfer is performed through the Direct Memory Access engine. `dlc_dma` (direct memory access instruction for deep learning computation) in conjunction with `tensor_slice` (tensor slicing function) loads input data from high-bandwidth memory into the input buffer with 128-byte alignment granularity, completely bypassing the cache (CMEM) level. The core transpose operation is implemented by `trans_to_128` (a 128-element transpose processing function), which calls the TPU-specific `m_transpose_push` (vector register transpose push instruction) to complete the dimensionality conversion from [C] to [C, 128] at the vector register level. This single instruction processes 8 × 128 elements, avoiding the complex index calculations and CMEM interactions found in the general path. Finally, the reassembled result is written back to high-bandwidth memory via `dlc_dma` (direct memory access instruction for deep learning computation). Throughout the process, `dlc_sync` (a deep learning computation synchronization instruction) ensures that the DMA operation is completed before proceeding to the next step, maintaining a strict execution order.
[0071] Optionally, the vector-level transpose reassembly is implemented using the `trans_to_128` function. This function internally calls the `m_transpose_push` instruction, leveraging the TPU's register-level transpose capability to complete the 8×128 element dimension transformation within a single clock cycle, avoiding the computational overhead of block-by-block concatenation via the `cat128` instruction in the general branch. Here, the vector-level transpose reassembly is implemented using `trans_to_128` (a 128-element transpose processing function), which internally calls `m_transpose_push` (a vector register transpose push instruction). This utilizes the TPU's register-level transpose capability to complete the 8×128 element dimension transformation within a single clock cycle, avoiding the computational overhead of block-by-block concatenation via the `cat128` (128-element vector concatenation instruction) in the general branch. The hardware acceleration mechanism fully leverages the special architecture of the TPU vector processing unit, simplifying the computational path in transpose reassembly scenarios.
[0072] It's worth noting that traditional methods require loading input data into cache blocks first, then concatenating and reassembling it using multiple `cat128` instructions, involving at least 24 clock cycles. In contrast, the embodiment of this application directly performs dimension mapping at the vector register level, achieving an equivalent operation with a single `m_transpose_push` instruction. This not only eliminates the storage overhead of intermediate data in the cache and reduces transpose latency, but also improves the throughput of key dimension transformation operations in the attention mechanism during actual model inference, while simultaneously reducing energy consumption and enhancing the processing efficiency of large models.
[0073] Further optionally, the triggering conditions for the transpose recombination scenario further include: the dimension interchange relationship between the input tensor and the target output tensor satisfies that the first dimension of the input tensor is equal to the second dimension of the target output tensor and the second dimension of the input tensor is equal to the first dimension of the target output tensor, wherein at least one dimension of the input tensor or the target output tensor is a single-dimensional structure with a preset threshold; and / or, the at least one dimension satisfies the preset alignment granularity constraint of the vector parallel computing unit; and / or, the data types of the input tensor and the target output tensor are consistent and belong to a preset supported data type set; and / or, the total number of elements in the input tensor is consistent with the dimension that satisfies the alignment granularity constraint, so as to eliminate the influence of redundant padding elements on the recombination result.
[0074] Understandably, the triggering conditions for transpose-based recombination scenarios are designed as a multi-dimensional constraint system to ensure that the hardware-accelerated path is activated only in optimal scenarios. Specifically, firstly, the dimension interchange relationship is verified, i.e., checking whether the first dimension of the input tensor is equal to the second dimension of the target output tensor (judged by the condition in->dim0 == out->dim1) and whether the second dimension of the input tensor is equal to the first dimension of the target output tensor (judged by the condition in->dim1 == out->dim0). Secondly, single-dimensional constraints are verified, requiring at least one dimension of either the input or target output tensor to be 1 (e.g., a single-dimensional structure check of in->dim1 == 1). Thirdly, alignment granularity verification is performed, using bitwise operations to check (in->dim0&127) == 0 (a 128-element alignment check) to ensure that the main dimension is an integer multiple of the vector processing unit granularity. Fourth, a data type consistency check is performed. The input data type is determined by `in->dtype == dlc_fp32`, and the output data type is determined by `out->dtype == dlc_fp32`. This confirms that both are 32-bit floating-point types and belong to the preset data type set supported by the TPU vector instruction set. Finally, a total element count verification is performed. The total number of elements in the input tensor is ensured to be exactly equal to the main dimension size by `in_numel == in->dim0`. This eliminates the impact of redundant elements introduced by dimension padding on the reconstruction accuracy.
[0075] Thus, this multi-condition judgment mechanism is completed by the scalar processing core within a short clock cycle. It replaces complex division operations with hardware-friendly bitwise operations (such as the bitwise AND operation in 128-element alignment check), reducing the condition judgment latency to the nanosecond level. For example, in the 128-layer Transformer architecture of the Qwen_72B large model, this trigger condition captures most of the [C,1] to [1,C] transpose operations, while keeping the false activation rate below 0.15%. This provides scenario adaptation for subsequent trans_to_128 (128-element transpose processing function) calls, ultimately improving the throughput of transpose operations.
[0076] Optionally, a pipelined collaborative mechanism is adopted to synchronously execute the steps of dividing the input tensor into blocks according to the capacity of the on-chip memory, and moving each block from high-bandwidth memory to the input block buffer through direct memory access, as well as calling the vector-level transpose reordering instruction or register-level transpose rearrangement mechanism of the vector parallel computing unit to arrange each block in the input block buffer according to the dimension of the target output tensor and write it into the output block buffer.
[0077] For example, multiple buffering is configured in on-chip memory for at least two input block buffers and at least two output block buffers; when data is transferred from high-bandwidth memory to on-chip memory in the first input block buffer, the control vector parallel computing unit performs data rearrangement and writes the data in the second input block buffer to the second output block buffer, and simultaneously writes the result in the first output block buffer back to high-bandwidth memory during the rearrangement write; through multiple buffering configuration and pipelined coordination mechanism, overlap is formed between the transfer stage and the computing stage, reducing waiting overhead and improving overall throughput.
[0078] In the above embodiments, the pipelined collaborative mechanism improves the efficiency of dimension reorganization through multi-buffer configuration and stage overlap. A four-buffer architecture is implemented in on-chip memory (VMEM): two input buffers, vmem_input_0 and vmem_input_1, and two output buffers, vmem_output_0 and vmem_output_1, are divided by vmem_allocator (on-chip memory allocator). Each buffer is of equal size and strictly aligned to the 128-element boundary. This configuration enables a triple pipeline: In the first stage (device 0), when `buffer_selector=0`, the first data block is moved from high-bandwidth memory to `vmem_input_0` via `dlc_dma` (direct memory access instruction for deep learning computation); in the second stage (device 1), `trans_to_128` (a 128-element transpose function) is called on the ready data in `vmem_input_1`. Internally, this function executes `m_transpose_push` (a vector register transpose push instruction) to complete dimensional rearrangement, and the result is written to `vmem_output_1`; in the third stage (device 0), the reassembled result in `vmem_output_0` is transferred in parallel to high-bandwidth memory via `dlc_dma_writeback` (a direct memory access instruction for writing back results). Buffer switching is automatically managed by the hardware state machine. When the first stage is completed, `buffer_selector` flips to 1, the original `vmem_input_0` becomes the write-back target, and new data begins to be loaded into `vmem_input_1`. Simultaneously, the second stage automatically switches to processing the new data in `vmem_input_0`. Key synchronization points are implemented using `pipeline_barrier`, which monitors the completion flags of the three stages and only allows buffer role switching when data dependencies are satisfied. For example, for the continuous reshape operation sequence in the Qwen_72B model, further optimization is achieved with a five-buffer configuration, introducing a dedicated prefetch buffer `vmem_prefetch`. The `prefetch_controller` preloads the data required for the next operation, reducing the pipeline bubble rate.
[0079] In the above embodiment of step S103, optionally, storage space is allocated in the on-chip memory for the input block buffer and the output block buffer, including: reserving a preset reserved area in the on-chip memory for system occupancy or other operator coordination, and aligning and trimming the available on-chip memory space according to a preset boundary alignment rule. The preset reserved area does not include a cache area, and all data processing is completed independently in the on-chip memory. Furthermore, within the available on-chip memory space after alignment and trimming, two consecutive buffer areas are divided according to the order of the input block buffer and the output block buffer, so that the starting address and capacity of the two consecutive buffer areas simultaneously meet the address alignment and length alignment requirements of direct memory access transfer, thereby reducing on-chip memory fragmentation and improving transfer efficiency. The data transfer process is completed directly in the on-chip memory area without passing through the cache.
[0080] In the embodiment of step S103, the on-chip memory space allocation process achieves optimal layout through multi-level precise control. First, vmem_reserve (on-chip memory reservation function) is called to reserve a fixed-size system reserved area (typically 64KB) at the VMEM starting address. This area is dedicated to storing hardware context and control state, explicitly excluding the cache (CMEM) area, ensuring that all data processing is completed independently in a closed loop within on-chip memory. Subsequently, alignment_calculator (alignment calculator) determines the starting boundary of the available space. The starting address for 16KB alignment is calculated using the formula usable_start = (reserved_end + ALIGN_BOUNDARY - 1)&(~(ALIGN_BOUNDARY - 1)), where ALIGN_BOUNDARY is dynamically determined by get_alignment_requirement (alignment requirement acquisition function) based on DMA engine characteristics. Next, capacity_trim (capacity trimmer) performs space trimming. `usable_size = (total_vmem_size - usable_start)&(~(MIN_BLOCK_SIZE - 1))` ensures the total capacity is an integer multiple of the minimum transfer block size (typically 128 bytes). Within the aligned, clipped space, `buffer_splitter` divides the space into two contiguous regions proportionally. `input_buffer_size = usable_size * 3 / 5`, `output_buffer_size = usable_size - input_buffer_size`, and buffer objects are created using `CxxTensor` (C++ tensor object constructor). `CxxTensorvmem_input(usable_start)` defines the input buffer, and `CxxTensor vmem_output(usable_start + input_buffer_size)` defines the output buffer. The critical alignment verification is performed by `dma_align_validator`, which checks two conditions: 1) the buffer start address satisfies `addr % DMA_MIN_ALIGNMENT == 0` (typically 128-byte alignment); 2) the buffer length satisfies `length % DMA_GRANULARITY == 0` (typically 128-element granularity). If verification fails, `buffer_realigner` automatically adjusts the boundary points until the requirements are met.Finally, `dma_config` (Direct Memory Access Configurator) sets the transfer parameters, specifying `vmem_input` and `vmem_output` as the source and destination addresses. Direct data transfer from HBM to VMEM is achieved through `dlc_dma` (Direct Memory Access Instruction for Deep Learning Computation), bypassing the cache throughout. The `memory_path_controller` forces the transfer path to HBM→VMEM instead of HBM→CMEM→VMEM. The entire allocation process is coordinated by a hardware state machine, and the generated buffer layout ensures efficient execution of subsequent dimension reorganization operations.
[0081] In the above embodiment of step S103, optionally, dividing the input tensor into blocks according to the capacity of on-chip memory includes: determining the maximum number of elements that can be moved in a single block based on the available capacity of on-chip memory, the data type bit width of the input tensor, and the preset vector length of the vector parallel computing unit. The maximum number of elements does not exceed the on-chip memory capacity, ensuring that the entire reassembly process does not require cache intervention. Furthermore, the input tensor is divided into blocks along the dimensional direction that satisfies the alignment granularity constraint, so that the length of each block satisfies the integer multiple constraint of the preset vector length. If the length of the tail block does not meet the integer multiple constraint, a preset tail block processing strategy is used to pad and move the tail block or perform rearrangement and write-back using a mask method to ensure the correctness of the reassembly result. The tail block processing is completed within on-chip memory and does not rely on the boundary processing capabilities of the cache.
[0082] In the above embodiments, the input tensor block partitioning process achieves efficient memory utilization through precise capacity calculation and dimension optimization. First, the block_size_calculator (block size calculation function) comprehensively evaluates the available on-chip memory capacity, the bit width of the input tensor data type (queried via get_bit_width (bit width acquisition function), and the preset vector length of the vector parallel computing unit (usually 128 elements), calculating the maximum number of elements in a single block: max_elements = (usable_vmem_size / 2) / element_size, where dividing by 2 ensures sufficient space in both the input and output buffers, and max_elements is adjusted to an integer multiple of the vector length, with the formula max_elements = (max_elements / VECTOR_LENGTH) * VECTOR_LENGTH. Subsequently, the dimension_selector analyzes the dimensions of each dimension of the input tensor and selects the dimension that satisfies the alignment granularity constraint as the block's main axis. This dimension must satisfy dim_size %VECTOR_LENGTH == 0, verified by the alignment_checker (alignment checker). Block splitting is performed by tensor_splitter, which splits the input tensor into equal-length blocks along a selected dimension, with each block's length being a strict integer multiple of VECTOR_LENGTH. For tail block processing, the tail_handler checks the remaining number of elements: tail_size = total_size % block_size. If tail_size > 0, one of two strategies is activated: when tail_size > VECTOR_LENGTH / 2, the padding_appender is called to pad the tail block to the full vector length, and an invalid region mask is generated by the zero_filler; when tail_size ≤ VECTOR_LENGTH / 2, the masked_rearranger is activated. This component calls vld_masked to read the tail data, combines it with mask_generator to create an accurate bitmap of valid elements, performs partial rearrangement by m_transpose_push_masked, and finally writes back the result by vst_masked, updating only the positions of valid elements.The entire tail block processing flow is executed in a closed loop within on-chip memory by the on-chip tail processor. A memory isolator ensures that cache regions are not accessed, and all masking operations are directly supported by hardware vector units, guaranteeing the numerical correctness of the reassembled results. This block-segmentation mechanism is coordinated by the block scheduler to generate an accurate block index table, providing fundamental control information for subsequent DMA transfers and dimension reassembly.
[0083] As an optional embodiment, in step S103, if it is determined that the input tensor and the target output tensor are consistent in terms of storage offset, padding method, continuity identifier, and linear access order, then it is determined that the current dimension reorganization operation belongs to the same layout direct memory access copy scenario, and it is determined that there is no need to perform vector calculation rearrangement, completely bypassing the data processing flow of on-chip memory and cache. Furthermore, direct memory access transmission parameters are generated according to the physical continuous interval of the input tensor or the transmission interval divided according to a preset step size, and the input tensor is copied from high-bandwidth memory to the high-bandwidth memory address corresponding to the target output tensor in a single or multiple direct memory access manner. The data transmission path does not pass through on-chip memory and cache, realizing direct transmission from high-bandwidth memory to high-bandwidth memory.
[0084] In the above embodiments, the implementation of direct copying within the same layout is achieved through precise metadata comparison and direct transmission mechanisms. First, the layout_comparator performs four-fold consistency verification: comparing the storage_offset field of the input / output tensors using the offset_matcher; checking whether the padding[0..4] parameters of each dimension are the same using the padding_validator; verifying whether the is_contiguous identifier is consistent using the contiguity_checker; and confirming that the linear access order (i.e., the relative proportions of the strides array) remains unchanged using the stride_analyzer. When all verifications pass, the direct_copy_decision activates the direct copying path within the same layout. Subsequently, the `dma_param_generator` (DMA parameter generator) generates transfer parameters based on the tensor layout characteristics: if the tensor is physically continuous, the `contiguous_range_calculator` (contiguous range calculator) is called to determine the starting address and length of a single transfer; if the tensor is not continuous but the stride is regular, the `strided_transfer_planner` (stride transfer planner) is started to decompose the transfer into multiple sub-intervals of equal stride length, each sub-interval meeting the 128-byte alignment requirement. Parameter configuration is completed by `hbm_dma_config` (high-bandwidth memory DMA configurator), which sets the source / destination address, transfer length, stride parameters, and the transfer mode flag `BYPASS_VMEM_CMEM=1`, forcing the data path to bypass on-chip memory and cache. Actual data transfer is executed through `dlc_dma_direct` (deep learning computation direct memory access instruction), a dedicated instruction that establishes a direct HBM-to-HBM data channel at the hardware level, with the memory controller directly scheduling the data flow. For large tensor segmented transfers, the multi_segment_coordinator manages multiple transfer sequences to ensure data integrity. Finally, the dma_sync_barrier polls the transfer completion flag to ensure that task resources are released only after all data segments have been copied. The entire process is automatically coordinated by a hardware state machine, achieving zero-computational-overhead dimension reorganization under the lightweight control of the scalar processing core.
[0085] As an optional embodiment, in step S103, if it is determined that at least one dimension of the input tensor and the target output tensor is equal to the first preset vector width, and the corresponding dimension of the other tensor is equal to the second preset vector width, and the remaining dimensions satisfy the alignment requirements of the vector parallel computing unit and the on-chip memory capacity constraints, then the current dimension reorganization operation is determined to be a vector rotation special case scenario involving preset vector width conversion. The dimension order of the input tensor and the target output tensor is normalized to ensure that the dimensions to be converted are located at preset vector processing dimension positions. This normalization is performed in on-chip memory, avoiding loading intermediate results into the cache. Input block buffers and output block buffers are allocated in on-chip memory, and the input tensor is moved to the input block buffer in blocks via direct memory access. Data movement does not pass through the cache, reducing memory-level interaction. The vector rotation, channel rearrangement, or component reorganization mechanism of the vector parallel computing unit is invoked to convert the data layout of the first preset vector width to the data layout of the second preset vector width, and write it to the output block buffer. The vector rotation operation is performed directly in the vector register, without needing to load data into the cache for reorganization. Finally, the results in the output block buffer are written back to the target output tensor via direct memory access.
[0086] In the above embodiments, the processing flow for vector width conversion scenarios is divided into three stages: accurate determination, data preparation, and vector conversion. First, the `width_switch_detector` performs dimension matching verification: it checks whether there exists a pair of corresponding dimensions in the input / output tensors, with dimensions equal to the first preset vector width (e.g., 40) and the second preset vector width (e.g., 320), respectively. Simultaneously, the `alignment_validator` confirms that all other dimensions meet the 128-element alignment requirement, and the `capacity_checker` verifies that the total number of elements does not exceed the on-chip memory capacity limit. When a match is successfully determined, the `dimension_normalizer` reorganizes the tensor dimension order and calls the `permute_dims` function to move the dimensions to be converted to the preset vector processing position (usually the last two dimensions). This operation only modifies the dimension index mapping table in the metadata and does not trigger actual data movement. Subsequently, the vmem_allocator (on-chip memory allocator) allocates input / output block buffers at a 3:1 ratio, ensuring 128-byte alignment of the starting address. During the data transfer phase, the dlc_dma_loader (deep learning computation DMA loader) transfers input tensor blocks from high-bandwidth memory to the input buffer, setting the BYPASS_CACHE=1 flag to bypass the cache. The core transformation operations are performed by the vector_swapper (vector swapper): for 40... 320 transformation, calling m_rotate (vector element rotation and rearrangement instruction) to set a fixed offset parameter of 8; for 16 64-channel transformation activates channel_shuffle to rearrange channel groups; for 8-channel transformation... The 32-step transformation initiates `component_repack` to adjust the data packing granularity. All vector operations are performed directly in the vector register, with `register_direct_path` ensuring data does not pass through the cache. Finally, `dlc_dma_writer` (the deep learning computation DMA write-back unit) transfers the output buffer data to the target tensor location, waiting for the transfer to complete via `dma_sync` (the DMA synchronization instruction). The `width_switch_scheduler` coordinates this process, ensuring seamless transitions between stages and achieving efficient vector width conversion.
[0087] Optionally, the vector rotation, channel rearrangement, or component reorganization mechanism is implemented collaboratively using the `m_rotate` instruction and the `cmp128` vector mask instruction. The `m_rotate` instruction adjusts the position of elements in the vector register, while the `cmp128` instruction filters invalid elements. The combination of these two achieves efficient dimensional layout transformation. For example, the collaborative mechanism of `m_rotate` (vector element rotation and rearrangement instruction) and `cmp128` (128-element boundary comparison instruction) achieves instruction-level optimization of dimensional layout transformation. `m_rotate` rearranges the positions of all elements in the vector register with a single instruction, adjusting the relative positions of 128 elements within one clock cycle using a hardware-level cyclic shift unit; while `cmp128` dynamically generates a 128-bit mask vector, accurately marking valid / invalid element regions and providing boundary protection for the rotation operation. The combination of these two forms an atomic operation of rotation and filtering, avoiding the cumbersome block decomposition and reorganization process in traditional methods. In contrast, the general branch requires at least 6-8 instruction sequences to achieve the same transformation using get128 (128-element vector extraction instruction) and cat128 (128-element vector concatenation instruction): first decomposing the original vector, then reorganizing it according to the new layout, and finally handling boundary conditions. The cooperative mechanism compresses this process into 2 instructions, reducing the number of computation instructions and eliminating the storage overhead of intermediate data in the cache.
[0088] In the above embodiment of step S103, optionally, the dimension order of the input tensor and the target output tensor is normalized, including: when it is determined that the dimension to be converted is not located at a preset vector processing dimension position, the dimension order of the input tensor or the target output tensor is converted to a standard dimension order that satisfies the internal processing granularity of the vector parallel computing unit, so that the dimension to be converted is aligned with the preset vector processing dimension. The conversion operation is completed in on-chip memory and does not involve data temporary storage in the cache. After completing the vector width conversion and writing back to high-bandwidth memory, the target output tensor is restored from the standard dimension order to the original dimension order of the target output tensor, and device synchronization is performed to ensure output consistency. Thus, the dimension restoration operation is achieved through metadata adjustment, without the need for data rearrangement through the cache.
[0089] In the above embodiments, dimension order normalization employs a metadata-driven lightweight transformation mechanism. First, the `dimension_position_checker` verifies whether the dimensions to be converted are located at the preset vector processing dimension position (usually the last two dimensions). If not, the `dimension_reorder` generates a dimension permutation mapping table, recording the correspondence between the original dimension index and the target position. Subsequently, the `metadata_adjuster` modifies the dimension order field in the tensor descriptor, moving the dimensions to be converted to the preset position. This operation only updates the metadata index mapping table and does not trigger actual data movement. When physical reordering is required (e.g., due to memory layout mismatch), the `lightweight_transposer` performs minimal data reorganization in on-chip memory: using `vld` (vector loading instruction) and `vst` (vector storage instruction) in conjunction with a pre-computed index offset table, it directly completes the dimension order adjustment within the input block buffer, setting the `CACHE_BYPASS=1` flag throughout to avoid cache intervention. After completing the vector width conversion, the `dimension_restorer` restores the dimension order of the target output tensor from the standard order to the original order. This operation also only modifies the dimension index and stride parameter in the metadata, updating the memory stride value of each dimension through the `stride_recalculator`. Finally, the `device_synchronizer` calls `dlc_sync` to wait for all processing units to complete their operations, ensuring consistency between metadata updates and physical data. The entire dimension normalization and restoration process is dominated by metadata operations, with physical data reassembly performed in a closed loop within on-chip memory only when necessary. This avoids the intermediate storage overhead of the cache and achieves zero-copy optimization for dimension order conversion.
[0090] As an optional embodiment, in step S103, if it is determined that the dimension reorganization operation is a single-dimensional expansion, merging, or insertion / deletion of dimensions with a preset threshold size on a continuous storage tensor, and it is determined that the dimension reorganization does not change the linear order of elements in physical storage, then the current dimension reorganization operation is determined to belong to the single-dimensional expansion and metadata remapping scenario. The scenario determination includes verifying that the data in the cache does not need to be rearranged. If it is determined that the linear order of elements is not changed, the movement and recalculation of the input tensor are stopped, completely avoiding redundant loading and reorganization operations of data in the cache. Tensor description information of the target output tensor is generated based on the tensor description information of the input tensor, wherein the tensor description information includes at least the number of dimensions, the size of each dimension, stride information, storage offset, and continuity identifier. The stride information and continuity identifier are recalculated according to the shape parameters of the target output tensor, and the target output tensor is set as a view output sharing the same physical storage area as the input tensor. The view output mechanism allows the two tensors to share the same physical storage, eliminating the need to create independent copies in the cache.
[0091] In this optional embodiment, the single-dimensional unrolling and metadata remapping scenario achieves zero-computational-overhead dimension reorganization through pure metadata operations. First, the dimension_analyzer identifies the operation type. When a single-dimensional unrolling (e.g., [128,1]→
[128] ), dimension merging (e.g., [16,8]→
[128] ), or single-dimensional insertion / deletion (e.g.,
[128] →[128,1]) is detected, the scenario determination process is activated. Subsequently, the linear_order_validator performs key checks. By comparing the linear index mapping function of the input / output tensors, it verifies whether index_in = index_out holds true, that is, the relative positions of the elements in physical storage remain unchanged. At the same time, the cache_rearrangement_checker confirms that the data in the cache does not need to be rearranged; if this verification fails, it falls back to the general processing path.
[0092] After successful verification, the `metadata_generator` constructs the target output tensor descriptor based on the input tensor's description information: copying the original storage pointer, data type, and other basic attributes, while adjusting the dimension count field. The `stride_recalculator` recalculates the stride array based on the target shape parameters. For single-dimensional deletion operations (e.g., deleting a dimension of size 1), it merges the stride values of adjacent dimensions. For single-dimensional insertion operations, it splits the corresponding dimension's stride and inserts a new stride value of 1. The `continuity_updater` re-evaluates the tensor's continuity flag, marking it as continuous when all dimension strides satisfy `stride[i] = stride[i+1] * dim[i+1]`. Finally, the `view_creator` sets the target output tensor as a shared view of the input tensors, using the `physical_storage_sharer` to make both point to the same physical memory region, setting the `is_view=true` flag, and recording the original tensor reference. The entire process is coordinated by metadata_only_scheduler (pure metadata scheduler), avoiding the overhead of data movement and computational rearrangement, and achieving ultimate optimization of dimension reorganization operations.
[0093] Further optionally, after completing the metadata remapping in step S103, it is determined whether to enter other special processing branches based on the dimensional characteristics of the remapped tensor. If the remapped tensor meets the specific dimensional conversion conditions, the corresponding dedicated acceleration function is called for efficient calculation.
[0094] Further optionally, after completing the tensor dimension reorganization within the on-chip storage range in step S103, the output result should be subjected to dimension restoration processing, and the status of each processing unit should be synchronized to ensure data consistency.
[0095] Understandably, in step S103, after a successful high-frequency special judgment pattern match, the tensor dimension reorganization is completed using hardware acceleration within the on-chip memory, supplemented by dimension restoration processing and processing unit state synchronization. Performing dimension reorganization within on-chip memory (VMEM) avoids redundant data movement between multiple storage levels, reducing memory access latency. For example, in the transpose scenario from [C,1] to [1,C], reorganization is completed directly in VMEM using vector register-level transpose instructions, reducing the number of CMEM interactions compared to the general path. Dimension restoration processing further optimizes the data layout, ensuring that the output tensor meets the memory access pattern requirements of downstream operators, avoiding additional rearrangement overhead in subsequent operations. When dimension reorganization involves multi-processing unit collaboration (e.g., device_id=0 is responsible for data loading, and device_id=1 is responsible for writing back the result), the state synchronization mechanism ensures consistent visibility of the reorganization result across all processing units through hardware semaphores and memory barrier instructions. This mechanism employs a lightweight two-stage synchronization protocol: first, a completion signal is broadcast via the on-chip interconnect network; then, the processing status registers of each unit are verified; and task resources are only released after all confirmations are received. This reduces the multi-core inconsistency error rate while minimizing synchronization overhead and overall execution time. On-chip reassembly combined with dimensionality restoration enables the recycling of storage resources. Specifically, the input block buffer is immediately reclaimed for use as the output buffer after data consumption. In 320 to 40-dimensional conversion scenarios, this mechanism improves effective data throughput while reducing VMEM fragmentation.
[0096] Step S104: When the current dimension reorganization operation does not match the high-frequency special judgment mode, a general reorganization execution method oriented towards hierarchical storage structure is adopted to complete the tensor dimension reorganization through parallel computing and collaborative processing.
[0097] As an optional embodiment, in step S104, when the minimum dimension size of the input tensor or the target output tensor is 1, the tensor with a minimum dimension size of 1 is converted into a one-dimensional layout to optimize memory access continuity; the general reorganization task is assigned to multiple tensor processing units, the first processing unit is responsible for loading and caching the input data, and the second processing unit is responsible for data reorganization and writing back the results; the first processing unit and the second processing unit are tensor processing units; the first processing unit loads the input tensor from high-bandwidth memory into on-chip memory in blocks, integrates the data elements in each data block through vector-level concatenation operations, and caches them in the cache; the second processing unit reads the data blocks from the cache, reorganizes the data into the dimensional format of the target output tensor through vector-level splitting operations, and temporarily stores it in on-chip memory; the temporarily stored reorganized data is written back from on-chip memory to the target output tensor in high-bandwidth memory; when dimensional preprocessing has been performed, the dimensional structure of the target output tensor is restored, and the states of the first processing unit and the second processing unit are synchronized to ensure data consistency.
[0098] Specifically, when step S104 determines that the current dimension reorganization operation does not match any high-frequency special judgment pattern, a general reorganization execution process for hierarchical storage structures is initiated. First, dimension preprocessing optimization is performed. When the lowest dimension size of the input tensor or target output tensor is detected to be 1 (e.g., shape [N,1] or [1,M]), the tensor is converted into an equivalent one-dimensional layout [N] or [M] using `reshape_view` (the dimension view reconstruction function), eliminating the memory access discontinuity caused by single dimensions and improving subsequent data loading efficiency. Then, the reorganization task is assigned to two cooperating tensor processing units. For example, `device_id=0` acts as the first processing unit, responsible for data loading and caching. `device_id=1` acts as the second processing unit, responsible for data reorganization and result writing back. Specifically, the first processing unit loads the input tensor from high-bandwidth memory (HBM) into on-chip memory (VMEM) in chunks using the `dlc_dma` (direct memory access instruction for deep learning computation). The chunk size is dynamically calculated based on the VMEM capacity as `chunk_size = min(VMEM_total_size / 2, 1024*1024)`. Next, it calls `cat128` (128-element vector concatenation instruction) to integrate the elements from each data block according to memory contiguity requirements, and then caches the integrated data in the cache (CMEM) using the `dlc_vld` (vector loading instruction). Simultaneously, the second processing unit reads data blocks from the cache using the `dlc_vst` (vector storage instruction), performs vector-level splitting using the `get128` (128-element vector extraction instruction), reassembles the data according to the dimension format of the target output tensor, and temporarily stores the intermediate results in the on-chip memory output buffer. After processing all data blocks, the second processing unit again calls `dlc_dma` (direct memory access instruction for deep learning computation) to write the reassembled results from on-chip memory back in batches to the target output tensor location in high-bandwidth memory. If dimensionality preprocessing was performed previously, the system calls `shape_restore` (dimensionality restoration function) to adjust the metadata of the target output tensor and restore its original multidimensional structure. Finally, `dlc_sync` (deep learning computation synchronization instruction) ensures that the states of the two processing units are completely synchronized. This instruction triggers a hardware barrier mechanism, forcing the completion of all pending memory operations and verifying data consistency flags to prevent computational errors caused by inconsistencies in the states between processing units.
[0099] Optionally, in the above steps, the first processing unit loads the input tensor from high-bandwidth memory into on-chip memory in blocks, and integrates and caches the data elements in each block into a high-speed cache through vector-level concatenation operations. This includes: configuring direct memory access parameters to ensure that the data transmission granularity is aligned with the vector processing width of the tensor processing unit; transferring data blocks from high-bandwidth memory to on-chip memory in batches through the direct memory access engine, optimizing the transmission granularity to match the storage bus width; reading data elements from on-chip memory into a vector register through vector loading instructions; recombining non-contiguous data fragments into logically contiguous data blocks in the vector register using vector concatenation instructions, and dynamically generating boundary mask bits using a vector mask control mechanism when the data block boundary involves a padding region, which is used to control the scope of the vector concatenation operation to filter invalid padding data; and caching logically contiguous data blocks into a high-speed cache through dynamic storage instructions.
[0100] For example, the transfer parameters are first set using `dlc_dma_config` (the direct memory access configuration function for deep learning computation) to ensure that the data transfer granularity is precisely aligned with the 128-element vector processing width of the TPU. This includes setting key parameters such as source / destination address offset, transfer stride, and block size. Then, `dlc_dma` (the direct memory access instruction for deep learning computation) is called to transfer data blocks in batches from high-bandwidth memory to on-chip memory. The transfer granularity is optimized to an integer multiple of 128 bytes, perfectly matching the storage bus width and achieving peak bandwidth utilization. After the data arrives at the on-chip memory, it is loaded using `vld` (the vector loading instruction) in conjunction with `BitWidthInfo`. <32> The `::AddrShift` (32-bit data address offset configuration) parameter efficiently reads data elements into an 8×128 vector register array. For non-contiguous memory layouts, the `cat128` (128-element vector concatenation instruction) is called to dynamically reassemble data fragments in the vector registers, merging previously scattered memory regions into logically contiguous data blocks. When processing data block boundaries and involving dimension padding regions, the `cmp128` (128-element boundary comparison instruction) generates a dynamic mask bitmap. This mask, through the `vmsk` (vector mask control instruction), precisely controls the scope of the `cat128` instruction, automatically filtering invalid padding data and ensuring that only valid elements participate in the reassembly calculation. Finally, the `fxcstDyn` (dynamic storage instruction) writes the logically contiguous data block from the vector registers to a specified location in the cache. This instruction automatically adjusts the storage range based on the actual effective size of the data block, avoiding writing to invalid regions.
[0101] Optionally, in the above steps, the second processing unit reads data blocks from the cache, reassembles the data into the dimensional format of the target output tensor through vector-level splitting operations, and temporarily stores it in on-chip memory. This includes: loading logically continuous data blocks cached by the first processing unit from the cache into a vector register; obtaining the logically continuous data blocks using vector loading instructions; calculating the mapping position of each element in the logically continuous data blocks in the target dimensional layout based on the dimensional size, dimensional arrangement order, and storage stride of the target output tensor; decomposing the logically continuous data blocks in the vector register into multiple sub-vectors using vector splitting instructions according to the mapping positions, and adjusting the relative positions of the elements in each sub-vector using vector rotation instructions; when performing a specific dimensional scaling conversion, setting the offset parameter of the vector rotation instructions to a preset value, performing element rearrangement using the mathematical relationship between dimensions, and obtaining the reassembled data blocks after position adjustment; and writing the reassembled data blocks into on-chip memory using vector storage instructions as intermediate reassembly results.
[0102] Continuing the example above, in the efficient data reorganization and write-back process executed by the second processing unit, the logically contiguous data block cached by the first processing unit is first loaded from the cache into an 8×128 vector register array using the vld (vector loading instruction). Subsequently, based on the dimension size array of the target output tensor, the dimension arrangement order, and the storage stride parameter, the scalar processing core calculates the linear mapping position of each element in the target dimension layout, generating a position offset table. According to this mapping table, the system calls get128 (128-element vector decomposition instruction) to decompose the contiguous data block in the vector register into multiple sub-vectors conforming to the target dimension granularity. For scenarios requiring element rearrangement, m_rotate (vector element rotation rearrangement instruction) adjusts the relative positions of elements within the sub-vectors according to the pre-calculated rotation offset. For example, when processing the dimension ratio conversion between 320 and 40, the offset parameter of m_rotate is fixed at 8 (320 / 40=8), utilizing the mathematical proportional relationship between dimensions to complete the element rearrangement in one step, avoiding complex index calculations. The reassembled data block is written to the on-chip memory output buffer via VST (vector storage instruction) to form an intermediate reassembly result.
[0103] Furthermore, in the above steps, writing the temporarily stored reassembled data from on-chip memory back to the target output tensor in high-bandwidth memory includes: after completing the reassembly processing of all logically contiguous data blocks, transferring the intermediate reassembly results in on-chip memory in batches to high-bandwidth memory through direct memory access to form the final target output tensor.
[0104] Continuing the example above, once all data blocks have been processed, `dlc_sync` (a deep learning computation synchronization instruction) is called to ensure the completion of the second processing unit's computation. Subsequently, `dlc_dma` (a deep learning computation direct memory access instruction) is used to batch transfer the complete reassembled results from on-chip memory to the target output tensor location in high-bandwidth memory. Thus, the batch DMA write-back operation reduces the final data transfer bandwidth, improving the execution efficiency and resource utilization of the general-purpose dimensional reassembly operation.
[0105] In another optional embodiment, in step S104, when the current dimension reorganization operation does not match the high-frequency special judgment mode, a general reorganization execution method oriented towards hierarchical storage structure is adopted. Tensor dimension reorganization is completed through parallel computing and collaborative processing, including: determining the block strategy and scheduling strategy for the input tensor and the target output tensor based on the on-chip memory capacity, high-bandwidth memory bandwidth, and alignment granularity of the vector parallel computing unit of the hierarchical storage structure; allocating storage space for the input block buffer and the output block buffer in on-chip memory, and establishing a mapping relationship between the block index and the global index based on the block strategy; performing loop processing on each block, and moving each block of the input tensor from the high-bandwidth memory to the input block buffer through direct memory access; calculating the address mapping between the input tensor and the target output tensor for each element in the current block based on the dimension size, dimension arrangement, stride information, and storage offset of the input tensor and the target output tensor, and having the vector parallel computing unit perform reading and writing in a vector parallel manner to complete the reorganization; and writing the reorganization result in the output block buffer back to the target output tensor in the high-bandwidth memory through direct memory access.
[0106] Understandably, in the above embodiments, it is assumed that the scheduler comprehensively analyzes the on-chip memory (VMEM) capacity (e.g., 3-6MB), the peak bandwidth of high-bandwidth memory (HBM) (e.g., 1.2TB / s), and the 128-element alignment granularity of the vector processing unit, and determines the optimal block parameters through the block_size_calculator (block size calculation function). This block_size_calculator function performs the calculation vmemlen = min(info->vmem_size / 4, 1024×1024×3)&(-128) to ensure that the block size neither exceeds the VMEM capacity limit nor is strictly aligned to the 128-element boundary. Subsequently, contiguous memory space is allocated in the VMEM using CxxTensor (C++ tensor object constructor). The `CxxTensor vmem_input(info->vmem_addr)` function defines the starting position of the input buffer, and the `CxxTensor vmem_output(vmem_input + vmemlen / 2)` function defines the output buffer. The ratio of the two is 1:1 to balance read and write requirements. At the same time, a block index mapping table `block_map` is established to record the global offset and local index relationship of each block. This mapping table takes into account the dimension stride and storage offset parameters to ensure the accuracy of index transformation.
[0107] Data processing employs a dual-core pipelined architecture. Based on this, device_id=0 (the first processing unit) is responsible for data loading. It uses dlc_dma (direct memory access instruction for deep learning computation) in conjunction with tensor_slice (tensor slicing function) to move input blocks from HBM to vmem_input, with the transfer parameters set to a 128-byte alignment granularity to match the memory bus width. device_id=1 (the second processing unit) is responsible for data reassembly. It reads preprocessed data from CMEM using fxcld (cache loading instruction), calls elemcpy_v2c (vector to cache copy function) combined with vld (vector loading instruction) and cat128 (128-element vector concatenation instruction) to integrate discontinuous data into logically contiguous blocks, and caches them to CMEM using fxcstDyn (dynamic storage instruction). In the core reorganization phase, based on the dimension size array (shape[0..4]) of the input / output tensors, the dimension arrangement order, and stride information, the address mapping function is calculated: new_idx = (old_idx / in_dim[i]) *out_dim[j] + (old_idx % in_dim[i]). The vector processing unit performs element-level rearrangement using get128 (a 128-element vector extraction instruction). Optionally, the read and write operations described in step S104 include processing the boundaries of filled or non-contiguous regions. For boundary regions, cmp128 (a 128-element boundary comparison instruction) dynamically generates a mask bitmap, and vmsk (a vector mask control instruction) filters the filled data to ensure that only valid elements participate in the calculation. The reorganization result is written to vmem_output using elemcpy_c2v (a cache-to-vector copy function) and then batch-transferred to HBM using the dlc_dma instruction.
[0108] Further, optionally, in step S104, synchronization or a pipelined transport and computation overlap mechanism is performed between adjacent blocks to improve the efficiency of general reorganization while ensuring data consistency.
[0109] Further optionally, in the above embodiments, establishing the mapping relationship between the block index and the global index based on the block strategy includes: constructing a mapping function from the coordinates of the output element to the coordinates of the input element based on the stride information and storage offset information of the input tensor and the target output tensor; generating the input read interval and output write interval for each block according to the mapping function; and performing read / write control on the boundary elements using preset boundary processing rules when the input read interval or output write interval contains a filled region, or when the input read interval or output write interval contains a non-contiguous region, so as to avoid crossing invalid storage areas and ensure the correctness of the reassembly result.
[0110] For example, firstly, based on the stride array and storage offset information of the input tensor and the target output tensor, a precise dimension mapping function is constructed using `coord_mapper` (a coordinate mapping generation function). This function parses the N-dimensional structure of the tensor, converting the global coordinates (out_coord[0..N-1]) of the output elements into the corresponding input coordinates (in_coord[0..N-1]), with the conversion formula: `in_coord[i]=(out_coord[0]*out_stride[0]+ ... +out_coord[N-1]*out_stride[N-1] +out_offset-in_offset) / in_stride[i]% in_dim[i]`. Subsequently, `block_range_calculator` (a block range calculation function) uses this mapping function, combined with the starting index and block size of the current block, to generate the input reading range [input_start, input_end] and output writing range [output_start, output_end] for each block. When these intervals are detected to cross dimensional padding regions or non-contiguous memory regions, the `boundary_handler` is activated. This controller performs triple boundary verification: First, the `padding_verifier` checks whether the elements of the input / output interval are within the valid data range, using the formula `elem_index < (dim_size - padding_size)`. Second, the `contiguous_checker` verifies whether the memory access sequence is continuous, identifying memory jumps caused by stride mismatch. Finally, the `mask_generator` dynamically generates a 128-bit mask vector to mark the positions of valid / invalid elements. A hierarchical processing strategy is used for boundary elements. When the input read interval contains a padding region, `vmsk_load` (mask load instruction) combined with `mask_vector` (mask vector) loads only valid elements, skipping padding data. When the output write interval crosses a non-contiguous region, `scatter_write` (hash write instruction) writes elements to the non-contiguous target positions according to a pre-computed offset table, rather than simply writing them continuously. At the physical implementation level, these boundary processing mechanisms are executed through hardware acceleration units. When is_boundary_block (boundary block identifier) is set, the TPU's scalar core automatically loads the preset boundary processing rule set and configures the mask register and address generation unit of the vector processing unit.
[0111] Optionally, in the above embodiments, cyclic processing is performed on each block, including: after each block is processed, updating the global offset and block count, and switching to the tail block processing stage when the number of remaining unprocessed elements is detected to be less than a preset block threshold. In the tail block processing stage, the transport parameters and rearrangement parameters of the tail block are generated based on the number of remaining unprocessed elements, and a fine-grained read / write strategy different from that of conventional block processing is used to complete the tail block reorganization, so as to reduce the resource waste caused by tail block processing while ensuring correctness.
[0112] Specifically, two key state variables are maintained: global_offset (global offset counter) records the total number of processed elements, and block_counter (block counter) tracks the serial number of the currently processed block. After each block is processed, global_offset += current_block_size is executed through offset_updater (offset update function), and the block count is incremented by block_counter++. Subsequently, remaining_checker (remaining element detection function) calculates the number of unprocessed elements remaining = total_elements - global_offset and compares it with the preset block threshold BLOCK_THRESHOLD (usually set to twice the width of the vector processing unit, i.e., 256 elements). When remaining < BLOCK_THRESHOLD, the system triggers the tail block processing flow. First, tail_param_generator (tail block parameter generator) dynamically calculates special transfer parameters based on the number of remaining elements, including tail_size = remaining, tail_align = (tail_size + 31)&-32 (32-byte alignment adjustment), and mask_bits = (1<<tail_size) - 1 (bitmask generation). Second, configure the DMA engine to use dlc_dma_tail (tail data direct memory access instruction), which accepts variable-length transfer parameters, avoiding resource waste caused by the fixed 128-byte granularity in regular blocks. In the data reorganization stage, vld_masked (masked vector load instruction) is called in combination with dynamic_mask (dynamic mask register) to only load valid elements and skip invalid data in the tail region. Subsequently, partial_transpose (partial transpose processing function) adjusts the rearrangement algorithm according to the actual size of the tail block, avoiding performing invalid calculations on non-existent elements. For tail blocks smaller than the vector width, scalar_fallback (scalar fallback processor) takes over the processing and completes the precise reorganization through a sequence of scalar instructions to ensure numerical correctness. In the write-back stage, vst_partial (partial vector store instruction) only writes to the valid data area to prevent overwriting adjacent memory.
[0113] The embodiment of this application improves the execution efficiency of the reshape operation, reduces the occupation of on-chip memory resources, helps avoid bandwidth resource waste, improves hardware utilization, provides efficient support for dynamic tensor operations in various machine learning models, and enables the tensor processing unit to show stronger adaptability and computing efficiency when facing complex model structures.
[0114] After introducing the method of the exemplary embodiments of this application, the following will describe a tensor dimension reorganization apparatus for a tensor processing unit according to an exemplary embodiment of this application. Figure 2 As shown, the device is applied to a tensor processing unit, which has a hierarchical storage structure and a vector parallel computing unit. The device 20 includes the following units: an input unit 201, used to verify whether the total number of elements in the input tensor and the target output tensor are equal to satisfy the basic constraints of tensor dimension reorganization; a judgment unit 202, used to perform mode determination on the current dimension reorganization operation based on the shape parameters and storage layout information of the input tensor and the target output tensor, as well as the hardware architecture characteristics of the tensor processing unit, to determine whether the current dimension reorganization operation matches a preset high-frequency special judgment mode; and a processing unit 203, used to, when the current dimension reorganization operation matches the high-frequency special judgment mode, use a hardware-accelerated execution method corresponding to the high-frequency special judgment mode to complete the tensor dimension reorganization within the on-chip storage range to obtain the target output tensor; when the current dimension reorganization operation does not match the high-frequency special judgment mode, use a general reorganization execution method oriented towards the hierarchical storage structure to complete the tensor dimension reorganization through parallel computing and collaborative processing to obtain the target output tensor. The above device can implement the steps described in the above method embodiments, and the specific implementation methods of each step will not be repeated here.
[0115] After introducing the methods and apparatus of the exemplary embodiments of this application, a terminal device of the exemplary embodiments of this application will be described next. The terminal device can implement the steps described in the above method embodiments, and the specific implementation of each step will not be repeated here.
[0116] After introducing the methods, apparatus, and terminal devices of exemplary embodiments of this application, the computer-readable storage medium of exemplary embodiments of this application will now be described. The computer-readable storage medium is an optical disc, on which a computer program (i.e., a program product) is stored. When the computer program is run by a processor, it implements the steps described in the above-described method embodiments. The specific implementation methods of each step will not be repeated here.
[0117] This application also relates to a chip system including at least one processor and an interface circuit. The processor includes a plurality of vector storage units. The processor is used to execute instruction and / or data interaction through the interface circuit, causing the chip system to execute the tensor dimension reorganization method for tensor processing units of any of the above embodiments.
[0118] In one possible implementation, the chip system may also directly include memory storing computer programs or computer instructions. For example, the memory may be volatile or non-volatile, or may include both. The non-volatile memory may be read-only memory (ROM), programmable read-only memory (PROM), erasable programmable read-only memory (EPROM), electrically erasable programmable read-only memory (EEPROM), or flash memory. The volatile memory may be random access memory (RAM), which serves as an external cache. By way of example, but not limitation, many forms of RAM are available, such as static random access memory (SRAM), dynamic random access memory (DRAM), synchronous dynamic random access memory (SDRAM), double data rate synchronous dynamic random access memory (DDR SDRAM), enhanced synchronous dynamic random access memory (ESDRAM), synchronous interconnected dynamic random access memory, and direct memory bus random access memory.
[0119] This application also relates to a processor, which includes a plurality of storage units for calling computer programs or computer instructions stored in the memory to cause the processor to execute the methods described in any of the above embodiments.
[0120] For example, in the embodiments of this application, the processor is an integrated circuit chip with signal processing capabilities. For instance, the processor may be an FPGA, a general-purpose processor, a DSP, an ASIC, or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, a SoC, a CPU, a network processor (NP), a microcontroller unit (MCU), a PLD, or other integrated chips, capable of implementing or executing the methods, steps, and logic block diagrams disclosed in the embodiments of this application. In one possible implementation, the embodiments of this application also provide a computer-readable storage medium storing program code, which, when executed on the computer, causes the computer to perform the above-described method embodiments.
[0121] It should be noted that the above-described embodiments are merely specific implementations of this application, used to illustrate the technical solutions of this application, and not to limit them. The protection scope of this application is not limited thereto. Although this application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that any person skilled in the art can still modify or easily conceive of changes to the technical solutions described in the foregoing embodiments, or make equivalent substitutions for some of the technical features, within the scope of the technology disclosed in this application. Therefore, the protection scope of this application should be determined by the scope of the claims.
Claims
1. A method for tensor dimension reorganization in a tensor processing unit, characterized in that, The tensor processing unit has a hierarchical storage structure and a vector parallel computation unit. The method includes: For the input tensor to be processed in the tensor processing unit, it is verified whether the total number of elements of the input tensor and the target output tensor are equal to meet the basic constraints of tensor dimension reorganization. Based on the shape parameters and storage layout information of the input and target output tensors, as well as the hardware architecture characteristics of the tensor processing unit, a mode determination is performed on the current dimension recombination operation to determine whether it matches a preset high-frequency special judgment mode. This includes: obtaining the dimensions and dimensional arrangement order of the input and target output tensors; determining whether the input and target output tensors satisfy a preset one-to-one dimension correspondence or dimension interchange relationship; and / or determining whether there exists a single-dimensional structure with a dimension size of a preset threshold in the input and / or target output tensors, and determining whether the shape constraint conditions of the high-frequency special judgment mode are satisfied based on the position of the single-dimensional structure in the tensor; and / or... The system acquires storage layout information for the input tensor and the target output tensor, the storage layout information including at least storage offset, padding method, and continuity identifier, and determines whether the input tensor and the target output tensor are physically continuous and consistent based on the storage layout information; and / or determines whether the data types of the input tensor and the target output tensor are consistent and whether they meet the data width and alignment conditions required by the vector instructions of the tensor processing unit; and / or determines whether they belong to the set of data types that can be directly recombined supported by the tensor processing unit; and / or, when at least one of the determination results meets the triggering condition corresponding to the high-frequency special judgment mode, the current dimension recombining operation is determined to match the high-frequency special judgment mode. When the current dimension recombination operation matches the high-frequency special judgment mode, the hardware acceleration execution method corresponding to the high-frequency special judgment mode is adopted to complete the tensor dimension recombination within the on-chip storage range and obtain the target output tensor. When the current dimension reorganization operation does not match the high-frequency special judgment mode, a general reorganization execution method oriented towards hierarchical storage structure is adopted. Tensor dimension reorganization is completed through parallel computing and collaborative processing to obtain the target output tensor.
2. The tensor dimension reorganization method for a tensor processing unit according to claim 1, characterized in that, The verification process checks whether the total number of elements in the input tensor and the target output tensor are equal to satisfy the basic constraints of tensor dimension reorganization, including: Obtain the number of dimensions and the size of each dimension from the tensor description information corresponding to the input tensor and the target output tensor respectively; Calculate the total number of elements in the input tensor based on the dimensions of each dimension of the input tensor; Calculate the total number of elements in the target output tensor based on the dimensions of each dimension of the target output tensor; If the total number of elements in the input tensor is the same as the total number of elements in the target output tensor, then the tensor dimension reorganization operation is performed.
3. The tensor dimension reorganization method for a tensor processing unit according to claim 1, characterized in that, When the current dimension reorganization operation matches the high-frequency special judgment mode, a hardware-accelerated execution method corresponding to the high-frequency special judgment mode is adopted to complete the tensor dimension reorganization within the on-chip storage range to obtain the target output tensor, including: Obtain the dimensional correspondence between the input tensor and the target output tensor, and determine that the current dimension recombination operation belongs to the transpose recombination scenario. The triggering conditions of the transpose recombination scenario include: the input tensor and the target output tensor are two-dimensional tensors and the first dimension and the second dimension of the input tensor and the target output tensor are interchanged, and at least one dimension is a single-dimensional structure with a preset threshold size. If the preset alignment granularity and data type constraints of the vector parallel computing unit of the tensor processing unit are satisfied, storage space is allocated in on-chip memory for the input block buffer and the output block buffer. The input tensor is divided into blocks according to the capacity of the on-chip memory, and each block is moved from high-bandwidth memory to the input block buffer through direct memory access. The vector-level transpose and reordering instruction or register-level transpose and rearrangement mechanism of the vector parallel computing unit is invoked to arrange each block in the input block buffer according to the dimension of the target output tensor and write it into the output block buffer. The reassembled results in the output block buffer are written back to the target output tensor in high-bandwidth memory via direct memory access.
4. The tensor dimension reorganization method for a tensor processing unit according to claim 3, characterized in that, The process of allocating storage space for the input block buffer and the output block buffer in on-chip memory to obtain the target output tensor includes: A preset reserved area is reserved in the on-chip memory for system use or other operator collaboration, and the available on-chip memory space is aligned and trimmed according to preset boundary alignment rules; Within the available on-chip storage space after alignment and trimming, two consecutive buffer regions are divided according to the order of the input block buffer and the output block buffer, so that the starting address and capacity of the two consecutive buffer regions simultaneously meet the address alignment and length alignment requirements of direct memory access transfer.
5. The tensor dimension reorganization method for a tensor processing unit according to claim 3, characterized in that, The step of dividing the input tensor into blocks according to the capacity of the on-chip memory includes: Based on the available capacity of on-chip memory, the data type bit width of the input tensor, and the preset vector length of the vector parallel computing unit, the maximum number of elements that can be moved in a single block is determined. The input tensor is divided into blocks along the dimensional direction that satisfies the alignment granularity constraint, such that the length of each block satisfies the constraint that the length of the preset vector is an integer multiple of the length. If the length of the tail block does not meet the integer multiple constraint, a preset tail block processing strategy is used to fill in and move the tail blocks or to perform rearrangement and write back in a masked manner.
6. The tensor dimension reorganization method for a tensor processing unit according to claim 1, characterized in that, When the current dimension reorganization operation matches the high-frequency special judgment mode, a hardware-accelerated execution method corresponding to the high-frequency special judgment mode is adopted to complete the tensor dimension reorganization within the on-chip storage range to obtain the target output tensor, including: If the input tensor and the target output tensor are determined to be consistent in terms of storage offset, padding method, continuity identifier, and linear access order, then the current dimension reorganization operation is determined to be a direct copy scenario of direct memory access in the same layout, and it is determined that no vector calculation rearrangement is required. Based on the physical continuous interval of the input tensor or the transmission interval divided by a preset step size, direct memory access transmission parameters are generated, and the input tensor is copied from high-bandwidth memory to the high-bandwidth memory address corresponding to the target output tensor in one or more direct memory accesses.
7. The tensor dimension reorganization method for a tensor processing unit according to claim 1, characterized in that, When the current dimension reorganization operation matches the high-frequency special judgment mode, a hardware-accelerated execution method corresponding to the high-frequency special judgment mode is adopted to complete the tensor dimension reorganization within the on-chip storage range to obtain the target output tensor, including: If it is determined that at least one dimension of the input tensor and the target output tensor is equal to the first preset vector width, and the corresponding dimension of the other tensor is equal to the second preset vector width, and the remaining dimensions satisfy the alignment requirements of the vector parallel computing unit and the on-chip storage capacity constraints, then If the current dimension recombination operation is determined to be a vector rotation special case of preset vector width conversion, the dimension order of the input tensor and the target output tensor is normalized so that the dimension to be converted is located at the preset vector processing dimension position. Input block buffers and output block buffers are allocated in on-chip memory, and input tensors are moved to the input block buffers in blocks through direct memory access; The vector rotation, channel rearrangement, or component recombination mechanism of the vector parallel computing unit is invoked to convert the data layout of the first preset vector width into the data layout of the second preset vector width, and then written into the output block buffer. The results in the output block buffer are written back to the target output tensor via direct memory access.
8. The tensor dimension reorganization method for a tensor processing unit according to claim 7, characterized in that, The normalization process for the dimension order of the input tensor and the target output tensor includes: When it is determined that the dimension to be converted is not located at the preset vector processing dimension position, the dimension order of the input tensor or the target output tensor is converted into a standard dimension order that satisfies the internal processing granularity of the vector parallel computing unit, so that the dimension to be converted is aligned with the preset vector processing dimension. After completing the vector width conversion and writing back to high-bandwidth memory, the target output tensor is restored from the standard dimensional order to the original dimensional order of the target output tensor, and device synchronization is performed to ensure output consistency.
9. The tensor dimension reorganization method for a tensor processing unit according to claim 1, characterized in that, When the current dimension reorganization operation matches the high-frequency special judgment mode, a hardware-accelerated execution method corresponding to the high-frequency special judgment mode is adopted to complete the tensor dimension reorganization within the on-chip storage range to obtain the target output tensor, including: If the dimension reorganization operation is determined to be a single-dimensional structure that expands, merges, or inserts or deletes dimensions of a continuous storage tensor with a preset threshold size, and if the dimension reorganization does not change the linear order of elements in physical storage, then the current dimension reorganization operation is determined to belong to the single-dimensional expansion and metadata remapping scenario. If the linear order of elements is not changed, stop performing the manipulation and recalculation of the input tensor. Tensor description information of the target output tensor is generated based on the tensor description information of the input tensor, wherein the tensor description information includes at least the number of dimensions, the size of each dimension, stride information, storage offset, and continuity identifier; The stride information and continuity identifier are recalculated based on the shape parameters of the target output tensor, and the target output tensor is set as a view output that shares the same physical storage area as the input tensor.
10. The tensor dimension reorganization method for a tensor processing unit according to claim 1, characterized in that, When the current dimension reorganization operation does not match the high-frequency special judgment mode, a general reorganization execution method oriented towards hierarchical storage structures is adopted. Tensor dimension reorganization is completed through parallel computation and collaborative processing to obtain the target output tensor, including: When the lowest dimension of the input tensor or the target output tensor is 1, convert the tensor with the lowest dimension of 1 into a one-dimensional layout. The general reorganization task is distributed to multiple tensor processing units. The first processing unit is responsible for loading and caching the input data, and the second processing unit is responsible for data reorganization and writing back the results. The first processing unit and the second processing unit are tensor processing units. The first processing unit loads the input tensor from high-bandwidth memory into on-chip memory in blocks, and integrates the data elements in each data block into a high-speed cache through vector-level concatenation operations. The second processing unit reads data blocks from the cache, reassembles the data into the dimension format of the target output tensor through vector-level splitting operations, and temporarily stores it in on-chip memory; Write the temporarily stored reassembled data back from on-chip memory to the target output tensor in high-bandwidth memory; When dimensionality preprocessing is performed, the dimensionality structure of the target output tensor is restored, and the states of the first processing unit and the second processing unit are synchronized to ensure data consistency.
11. The tensor dimension reorganization method for a tensor processing unit according to claim 10, characterized in that, The process of loading the input tensor from high-bandwidth memory into on-chip memory in blocks by the first processing unit, and integrating and caching the data elements in each block into a high-speed cache through vector-level concatenation operations includes: Configure direct memory access parameters to ensure that the data transfer granularity is aligned with the vector processing width of the tensor processing unit; The Direct Memory Access Engine (DME) enables bulk transfer of data blocks from high-bandwidth memory to on-chip memory, optimizing the transfer granularity to match the storage bus width. Data elements are read from on-chip memory into a vector register using vector load instructions; The vector concatenation instruction is used to reassemble non-contiguous data segments into logically contiguous data blocks in the vector register. When the boundary of the data block involves the padding area, the vector mask control mechanism dynamically generates boundary mask bits to control the scope of the vector concatenation operation and filter out invalid padding data. Logically contiguous data blocks are cached in the cache using dynamic storage instructions.
12. The tensor dimension reorganization method for a tensor processing unit according to claim 10, characterized in that, The step of reading data blocks from the cache by the second processing unit, reorganizing the data into the dimension format of the target output tensor through vector-level splitting operations, and temporarily storing it in on-chip memory includes: Load the logically contiguous data block cached by the first processing unit from the cache into the vector register, and use a vector load instruction to obtain the logically contiguous data block; Based on the dimension size, dimension arrangement order and storage stride of the target output tensor, calculate the mapping position of each element in the logically continuous data block in the target dimension layout; Based on the mapping position, the logically continuous data block in the vector register is decomposed into multiple sub-vectors using a vector splitting instruction, and the relative positions of the elements in each sub-vector are adjusted using a vector rotation instruction. When performing a specific dimension scaling conversion, the offset parameter of the vector rotation command is set to a preset value, and the elements are rearranged using the mathematical relationship between dimensions to obtain the recombined data block after position adjustment; The reconstructed data blocks are written to on-chip memory using vector storage instructions as intermediate reconstructing results.
13. A tensor dimension reorganization device for a tensor processing unit, characterized in that, The device is applied to a tensor processing unit, which has a hierarchical storage structure and a vector parallel computation unit. The device includes the following units: The input unit is used to verify whether the total number of elements of the input tensor and the target output tensor are equal to the input tensor to be processed, so as to meet the basic constraints of tensor dimension reorganization. The judgment unit is used to perform mode determination on the current dimension recombination operation based on the shape parameters and storage layout information of the input tensor and the target output tensor, as well as the hardware architecture characteristics of the tensor processing unit, so as to determine whether the current dimension recombination operation matches the preset high-frequency special judgment mode. The judgment unit is specifically used for: acquiring the dimensions and dimensional arrangement order of the input tensor and the target output tensor; determining whether the input tensor and the target output tensor satisfy a preset one-to-one correspondence or dimension interchange relationship; and / or determining whether there is a single-dimensional structure with a dimension size of a preset threshold in the input tensor and / or the target output tensor, and determining whether the shape constraint conditions of the high-frequency special judgment mode are satisfied based on the position of the single-dimensional structure in the tensor; and / or acquiring the storage layout information of the input tensor and the target output tensor, wherein the storage layout information includes at least storage offset, filling method and continuity identifier, and determining whether the input tensor and the target output tensor are physically continuous and consistent based on the storage layout information; and / or determining whether the data types of the input tensor and the target output tensor are consistent, and whether they meet the data width and alignment conditions required by the vector instruction of the tensor processing unit; and / or determining whether it belongs to the set of data types that can be directly recombined supported by the tensor processing unit; and / or, when at least one judgment result satisfies the triggering condition corresponding to the high-frequency special judgment mode, determining the current dimension recombining operation as matching the high-frequency special judgment mode; The processing unit is used to, when the current dimension reorganization operation matches the high-frequency special judgment mode, adopt a hardware-accelerated execution method corresponding to the high-frequency special judgment mode to complete the tensor dimension reorganization within the on-chip storage range and obtain the target output tensor; when the current dimension reorganization operation does not match the high-frequency special judgment mode, adopt a general reorganization execution method oriented towards hierarchical storage structure, and complete the tensor dimension reorganization through parallel computing and collaborative processing to obtain the target output tensor.
14. A chip, characterized in that, The chip stores a computer software program, which, when executed, implements the tensor dimension reorganization method for a tensor processing unit as described in any one of claims 1-12.