Optimal control Riccati equation solving method based on TPU heterogeneous parallelism

Abstract

The invention relates to the technical field of robot motion control and heterogeneous calculation acceleration, in particular to an optimal control Riccati equation solving method based on TPU heterogeneous parallelism, which is characterized in that heterogeneous division of labor is carried out on a host side and a TPU equipment side, the host side is only responsible for input issuing and result recovery, and the equipment side is used for managing and controlling an iterative process and operator scheduling by a coprocessor; the multiple neural network processing units execute intensive linear algebraic calculation in parallel; designing a general matrix multiplication parallel acceleration operator and a parallel Gaussian elimination inversion operator, optimizing a storage layout of an on-chip local memory and a working area reuse strategy, completing Riccati equation solution based on Hamiltonian matrix sign function iteration, recovering a stable solution matrix and calculating an optimal feedback gain; for a multi-independent system scene, a system dimension parallel or batch rotation strategy is adopted to improve the batch processing throughput rate. The method solves the problems that in the prior art, intermediate data carrying overhead is large, and parallel support of key links is insufficient.

Description

Technical Field

[0001] This invention relates to the field of robot motion control and heterogeneous computing acceleration technology, and particularly to a method for solving the Riccati equation for optimal control based on TPU heterogeneous parallelism. Addressing the numerical solution task of the continuous-time algebraic Riccati equation, this method employs heterogeneous parallelism to efficiently calculate the stable solution P and the optimal feedback gain K of the Riccati equation. It can be applied to solving optimal control algorithms such as linear quadratic regulation (LQR) and model predictive control (MPC). Background Technology

[0002] The continuous-time algebraic Riccati equations have wide applications in robot optimal control, model predictive control, state estimation, and trajectory tracking. Their solution matrix P is used to construct the optimal feedback gain K, achieving a trade-off between control system stability and performance. Common engineering methods for solving the Riccati equations include the stable subspace method based on Schur decomposition and eigenvalue rearrangement, the method based on Hamiltonian matrix-invariant subspaces, and numerical methods based on matrix sign function iteration. These methods are relatively mature on general-purpose processors (CPUs) or mathematical software environments such as Matlab and Eigen, but they still have many technical shortcomings in higher-dimensional systems or control applications requiring frequent repetitive solutions. 1. The general-purpose CPU / Matlab / Eigen solution suffers from high computational time and limited parallelism. The Riccati solution process typically involves matrix multiplication, matrix inversion or linear equation solving, projection partitioning, and convergence determination. It involves both rule calculations for multiplicative cryptography and control flow operations such as pivot selection, row swapping, and stopping criterion determination. For scenarios with high dimensionality or requiring repeated solutions in each control cycle, general-purpose CPU implementations are often limited by single-core throughput and memory bandwidth, making it difficult to meet the requirements for higher control frequencies and lower latency in terms of end-to-end time consumption.

[0003] 2. Conventional GPU acceleration solutions incur additional overhead and system constraints in small- to medium-scale matrix scenarios. The matrix dimensions commonly found in robot control are typically small to medium (e.g., n < 50). While GPU solutions offer high parallelism, their actual deployment is often affected by various factors such as kernel startup, communication overhead, power consumption and heat dissipation, and system integration complexity. This leads to unstable overall performance and makes it difficult to directly meet the combined requirements of real-time performance and engineering constraints in embedded / edge controllers.

[0004] 3. Existing TPU solutions lack sufficient support for the linear algebra operations required by control algorithms, easily leading to round-trip overhead between acceleration unit computation and general-purpose kernel inversion. Current TPU applications are mostly concentrated on inference / training operators such as convolution and matrix multiplication. They often lack efficient hardware operators to support essential steps in the control solution process, such as matrix inversion or linear equation solving. A common engineering practice is to roll back these steps to the general-purpose processing kernel. This introduces frequent transfer of intermediate matrices between the host and device. Repeated data reads and writes during iterations make external memory bandwidth a bottleneck, and it also makes it difficult to fully utilize the aligned layout and reuse advantages of on-chip local memory, thus limiting overall performance.

[0005] In summary, existing technologies for solving the Riccati equations on the TPU platform still lack an end-to-end solution mapping method that matches its computational / storage characteristics. In particular, there is a lack of implementation paths adapted to the parallel execution of acceleration units in key bottleneck links such as matrix inversion (or solving linear equations), making it difficult to balance real-time performance, bandwidth efficiency, and engineering deployability. Summary of the Invention

[0006] The purpose of this invention is to overcome the shortcomings of the prior art and provide a method for solving the Riccati equation of optimal control based on TPU heterogeneous parallelism. By reasonably dividing the control flow and dense linear algebra calculations, key operations such as matrix multiplication and matrix inversion in the iteration process can be efficiently executed by making full use of the TPU's on-chip resources. This solves the problems of large overhead of intermediate data round trips and external memory access and insufficient end-to-end computation efficiency when performing optimal control solutions on this type of platform in the prior art.

[0007] The above objectives are achieved through the following technical solutions: A method for solving the optimal control Riccati equation based on TPU heterogeneous parallelism includes the following steps: S1: On the host side, the system matrices A, B, Q, and R are organized into continuous data in row-major order and written to the device-side system memory of the Tensor Processing Unit (TPU) in one go. The Riccati equation solving kernel is then started, and after the calculation is completed, the stable solution matrix is ​​read back and output. With the optimal feedback gain matrix ; S2: On the device side, the coprocessor executes iterative process control and operator scheduling, and moves the input matrix from the system memory to the on-chip local memory through Global Direct Memory Access (GDMA), and completes data alignment and layout preparation; S3: Construct a Hamiltonian matrix in the on-chip local memory based on the input matrices A, B, Q, and R. Using the Frobenius norm pair Perform scaling preprocessing to obtain initial values ​​for iteration. ,in ;in, The Hamiltonian matrix is ​​constructed based on the system matrix. Describe the Frobenius norm. for The scaling factor is calculated using the Frobenius norm. This is the initial matrix for the symbolic function iteration; S4: Perform Newton-type iterations based on matrix sign functions. In each iteration, call the parallel Gaussian elimination inverse operator to calculate... The inverse matrix, according to the iterative update rule get The coprocessor calculates and determines whether the convergence criterion is met; if it is, the iteration exits. For the first The iteration matrix of the nth iteration. For the first The iteration matrix of the nth iteration. for The inverse matrix, The number of iterations ( ≥0); S5: After the iteration converges, construct the iterative result of the symbolic function. And based on this, a projection matrix is ​​formed. and its blocks , Solving linear equations by calling parallel Gaussian elimination inversion operator and general matrix multiplication operator. To restore the stable solution matrix ;in, Hamiltonian matrix The approximate value of the sign function, Represents matrix symbolic functions, Let be the projection matrix. To and Identity matrices of the same order This is the top-left block matrix of the projection matrix. This is the lower left block matrix of the projection matrix; S6: Based on optimal control relationship The optimal feedback gain matrix is ​​calculated using a general matrix multiplication operator and by calling a parallel Gaussian elimination inversion operator when needed. ,Will , Write back to the device-side system memory; among which, Input weight matrix The inverse matrix, For the input matrix The transpose of .

[0008] Preferably, in step S2, the matrices A, B, Q, R, and the iteration matrix are stored in the on-chip local memory. The augmented matrix buffer and projection block buffer of the inversion work area are partitioned and planned. Each matrix is ​​organized and stored using a uniform alignment row stride, and partitioned and reused according to the storage bank / buffer.

[0009] Preferably, in step S3, the Hamiltonian matrix H is constructed, and in step S4, iterative updates are performed. Step S5: Block Calculation of Projection Matrix and Stable Solution Recovery, optimal feedback gain in step S6 All matrix multiplications during the computation process are uniformly performed by the General Matrix Multiplication (GEMM) operator on the Neural Processing Unit (NPU) side of the Tensor Processor (TPU).

[0010] Preferably, the implementation process of the parallel Gaussian elimination inversion operator in steps S4 and S5 is as follows: the matrix to be inverted is stored in on-chip local memory. Constructing augmented matrices ,in To and For identity matrices of the same order, Gaussian elimination transforms the left half of the augmented matrix into an identity matrix, and the right half becomes the inverse of the original matrix M. .

[0011] Preferably, in the Gaussian elimination process... During the elimination process, the coprocessor selects a partial pivot selection strategy, choosing the element with the larger absolute value as the pivot, and performs row swapping when necessary. Row swapping is implemented in local memory by moving and overwriting data in whole row vectors. The pivot selection, row swapping decision and elimination step control are all handled by the coprocessor.

[0012] Preferably, the augmented matrix The matrix is ​​divided into several column-wise sub-blocks, with the sub-block width adapted to the vector parallelism width of the neural network processing unit (NPU). These column-wise sub-blocks are mapped to the local memory of multiple NPUs, with some NPUs storing the original matrix. The columns are partitioned, and the remaining neural network processing units (NPUs) store unit arrays. The column slices are divided into columns. Within the same elimination step, each neural network processing unit (NPU) performs vectorized operations of row normalization and row elimination update in parallel on different column slices of the same row based on the tensor multiplication constant and tensor subtraction instructions provided by the tensor processor (TPU) operator library. By setting the shape and stride of the tensor, each column slice participates in the row operation according to the unified address calculation rules.

[0013] Preferably, when solving the Riccati equations for multiple independent systems, the host side will handle the matrices of the N systems. The system is batch-packaged according to its index and delivered to the device-side system memory in one go. The device allocates or reuses on-chip local memory partitions for multiple systems, mapping these systems to multiple Neural Processing Units (NPUs) or groups of NPUs for parallel execution of symbolic function iteration, projection block construction, and... , Recovery; when the number of systems exceeds the available parallel resources, a batch rotation method is used to perform fully block-based parallel solution of multiple systems.

[0014] Preferably, the input matrices A, B, Q, and R and the output matrix... , The data type is single-precision floating-point (FP32), half-precision floating-point (FP16), or brain half-precision floating-point (BF16), depending on the operator support of the Tensor Processor (TPU) device.

[0015] Preferably, the convergence criterion in step S4 is calculated as follows: ,in The preset convergence threshold, It is the Frobenius norm.

[0016] Preferably, the on-chip local memory working area is multiplexed using a time-division multiplexing method, and is used to store the iteration matrix during the iteration phase. The storage area for the augmented matrix and its temporary results is reused during the inversion phase to store the augmented matrix and intermediate quantities of elimination. The reuse switching timing is uniformly managed by the coprocessor, and data conflicts are avoided by clearly defining the data lifecycle and reading / writing boundary control.

[0017] The present invention provides a TPU-based heterogeneous parallel method for solving the Riccati equations of optimal control. Compared with existing solutions based on general-purpose processors (CPU / Matlab / Eigen) or conventional discrete acceleration (general-purpose kernels that only perform matrix multiplication and inversion backoff), this invention organizes the complete process of solving the Riccati equations on the TPU device and heterogeneously divides the control flow and dense linear algebra operations. This allows key calculations such as matrix multiplication and matrix inversion required for iterative updates to be efficiently completed on on-chip resources, thereby reducing the dependence of the solution process on the host side and external memory bandwidth. This improves the efficiency of solving the solution matrix P and feedback gain K in optimal control and enhances the engineering deployability. The specific effects are as follows: 1. End-to-end device-side execution and on-chip storage collaboration: This invention unifies the iterative computation and intermediate matrix processing of the Riccati solution on the device side, while the host side only handles input delivery and result retrieval. Simultaneously, key intermediate quantities such as the iterative matrix are preferentially stored in local memory, interacting with system memory only through GDMA during initialization and final write-back phases. Compared to existing technologies that frequently trigger host-side and device-side transmissions or repeatedly access external memory during iteration, this invention reduces the overhead caused by intermediate data movement and external memory bandwidth usage.

[0018] 2. Heterogeneous Division of Labor and Multi-NPU Parallel Computation: In this invention, the device-side coprocessor is responsible for control flow tasks such as iteration organization, convergence determination, and operator scheduling. Multiple NPUs perform intensive linear algebra core calculations such as matrix multiplication and Gaussian elimination inversion, thereby avoiding the performance bottleneck caused by general-purpose processing cores handling large-scale matrix operations and improving the overall efficiency of end-to-end solution. Based on test data, the TPU's tensor processing unit provides dedicated optimization capabilities for multiplication-intensive set computations in general matrix multiplication. Through on-chip parallel computing units and instruction-level vectorized execution mechanisms, it can improve the throughput of matrix multiplication in small-to-medium scale linear algebra operations, thus providing stable computational support for matrix multiplication in Riccati iterations. Meanwhile, in the matrix inversion (or equivalent linear equation solving) stage, which is more likely to become a bottleneck in the Riccati solution, this invention utilizes multiple NPUs to perform Gaussian elimination in parallel, which can further reduce the time consumption of key steps. For example, in the inversion test of a 48×48 scale, the multi-NPU parallelism reduced the time from 1687 microseconds for a single NPU to 1100 microseconds, and achieved speed improvements of approximately 5.0 times and 5.6 times compared to the CPU baseline of 5491 microseconds and the Eigen baseline of 6192 microseconds, respectively.

[0019] 3. Alignment Layout and Working Area Reuse: This invention adopts a unified stride alignment and layout convention in local memory, and performs bank / buffer partitioning planning for input matrices, iteration matrices, augmented matrices, and projection blocks. The coprocessor uniformly manages the working area reuse and switching timing, reducing repeated allocation and release operations and data rearrangement operations. This allows matrix multiplication and inversion elimination operators to be read and written directly under a consistent data organization method, improving on-chip resource utilization efficiency and enhancing adaptability to problems of different scales.

[0020] 4. Engineering integration and reusability: This invention abstracts the key computational steps in the Riccati solution process into two core operators: matrix multiplication and matrix inversion. It provides a unified data layout, parameterized interface, and timing / error output mechanism, which facilitates modular integration and reuse in dedicated operator libraries or control system software stacks. At the same time, it supports comparison and verification with consistent testing standards.

[0021] 5. Numerical stability assurance: Through partial principal component selection strategy and Hamiltonian matrix scaling preprocessing, the relative error between the solution and the Matlab icare reference solution meets engineering requirements, ensuring reliability in control system applications.

[0022] The evaluation method and judgment criteria of this invention are as follows: Firstly, the performance statistics break down Riccati's end-to-end time into overall kernel time and device-side computation time, which is used to locate the main sources of time consumption and bottlenecks. Secondly, the correctness and numerical stability are verified by comparison with a reference implementation. Matrix multiplication is determined by the element-wise error threshold, and matrix inversion is verified by the consistency between the element-wise difference and the maximum absolute error threshold (or norm error). Third, the Riccati prototype evaluation uses the relative error index of the solution matrix P and the feedback gain K (relative to the Matlabicare solver reference solution) and the above timing results as a comprehensive evaluation criterion, and compares the consistency with benchmark schemes such as CPU / Eigen / Matlab. Attached Figure Description

[0023] Figure 1 This is a schematic diagram of the TPU heterogeneous computing architecture used for solving the Riccati equation (CARE) in the optimal control Riccati equation solving method based on TPU heterogeneous parallelism described in this invention. Figure 2 This is a flowchart of the solution process for the Riccati equation (CARE) based on the Hamiltonian matrix symbolic function iteration method in the optimal control Riccati equation solution method based on TPU heterogeneous parallelism described in this invention. Figure 3 This is a schematic diagram of the parallel Gauss-Jordan inverse operator process and the augmented matrix column partitioning multi-NPU mapping in the optimal control Riccati equation solution method based on TPU heterogeneous parallelism described in this invention. Detailed Implementation

[0024] The present invention will now be described in further detail with reference to the accompanying drawings and embodiments. The described embodiments are merely some, not all, of the embodiments of the present invention. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0025] This solution provides an optimal control Riccati equation solution based on TPU heterogeneous parallelism. Utilizing a TPU heterogeneous parallel architecture, it achieves efficient solution of the Riccati equation through the division of labor between the host and device, parallel design of core operators, optimization of memory layout, and multi-system parallel scheduling. The details are as follows: (a) Overall architecture and end-to-end execution mode, such as Figure 1 As shown 1. The host side is only responsible for sending input data and collecting results. The input matrices A, B, Q, and R are expanded row-majorly and written into the TPU global memory at once, and the device-side solving kernel is started. During the solving process, the host side does not participate in numerical calculations to avoid frequent matrix backhauls during iteration.

[0026] 2. The device side adopts a collaborative working mode of coprocessor and acceleration unit. The coprocessor (such as A53-Lite) is responsible for iterative process control and operator scheduling, including iterative loops, convergence determination, principal component strategy, work area management, and data transfer timing; the NPU computing instructions are responsible for executing high-throughput matrix multiplication and parallel Gaussian elimination inversion operators, and completing key linear algebra calculations such as multiplication encryption set and elimination update in the Riccati solution process.

[0027] 3. System memory and local memory collaboration and data resident status. On the device side, A, B, Q, R, iterative matrix, projection block and working area data should be kept in local memory as much as possible. The matrix storage should be organized according to the row step size of the alignment rule to reduce the bandwidth loss caused by non-aligned access. The same row step size convention can be directly reused for subsequent matrix multiplication and inversion, thereby reducing format conversion and data rearrangement.

[0028] (II) Mathematical process for solving the Riccati equation (based on Hamiltonian matrix symbolic function iteration), such as Figure 2 As shown For the continuous-time algebraic Riccati equations, the stable solution matrix Satisfies the following matrix equation: , The specific solution process is as follows: 1. Construct the Hamiltonian matrix and perform scaling preprocessing: Construct the Hamiltonian matrix based on the inputs A, B, Q, and R. , The expression is: , To improve the numerical stability and convergence of subsequent symbol function iterations, Frobenius norm pairs are used. Scaling is performed, and the scaling factor is calculated. To obtain the initial value of the iteration .

[0029] 2. Iterative update using Newton-type symbolic functions: The iterative formula is as follows The calculation of the inverse matrix is ​​accomplished by the parallel Gaussian elimination inverse operator designed in this invention.

[0030] 3. Convergence determination and output of the symbolic function iteration results: The coprocessor determines the convergence based on the convergence criteria. Determine if the iteration has ended; if the criterion is met, output the iteration result. ,in The preset convergence threshold is used; 4. Recovering the stable solution matrix based on projection relationships :based on Constructing the projection matrix And according to the state / conjugate variable space Divide into blocks to obtain By solving the system of linear equations Obtain the stable solution matrix .

[0031] 5. Calculate the optimal feedback gain matrix According to the optimal control relationship Calculate the feedback gain K, where matrix multiplication is accelerated using the general matrix multiplication operator, and matrix inversion is solved by calling the parallel Gaussian elimination inversion operator.

[0032] (III) General Matrix Multiplication Parallel Accelerator Operator The matrix multiplication operations involved in iterative updates, projection restoration, and feedback gain K calculation are uniformly mapped to a general matrix multiplication operator, which is then executed by the matrix multiplication calculation instructions on the NPU side. This operator adopts an output-oriented two-dimensional block strategy, dividing the output matrix into several sub-blocks along the row and column directions and distributing them to different NPUs. The multiplication and addition calculations of each sub-block are completed synchronously, and the piecewise accumulation is performed on the shared K dimension, thereby achieving parallel execution of matrix multiplication and improving throughput.

[0033] To reduce the overhead of global memory access and data rearrangement, input fragments, weighted fragments, and output fragments reside simultaneously in local memory. A unified storage convention of row-major order and alignment step size is adopted, and sub-block location is completed through fixed address calculation rules. For fragment size selection, the row and column fragment sizes are adaptively determined based on the local memory capacity and the number of NPUs, and the column fragment size is aligned to an integer multiple of the number of NPUs to achieve load balancing. Furthermore, during K-dimensional traversal, double buffering and a "prefetch-compute-write-back" pipeline are employed, allowing data movement and multiply-accumulate calculations to overlap, thereby reducing memory access wait times and maintaining continuous matrix multiplication calculations.

[0034] (iv) Parallel Gaussian elimination inverse operator, such as Figure 3 As shown 1. Basic process of constructing and inverting augmented matrices: [Following the steps of constructing and inverting the matrix...] Construct the augmented matrix in local memory (A same-order identity matrix is ​​concatenated to the right of the original matrix). Gaussian elimination is used to transform the left half of the augmented matrix into an identity matrix, and the right half is the inverse matrix of the original matrix. Taking a 48×48 single-precision floating-point matrix as an example, contiguous space is allocated in local memory for the 48×96 augmented matrix, and all elimination steps are performed on-chip.

[0035] 2. Pivoting Strategy and Row Swapping (Numerical Stability): To reduce the risk of numerical amplification caused by excessively small pivots, the coprocessor employs a partial pivoting strategy. In the k-th elimination step, the pivot element with the larger absolute value is selected as the pivot, and row swapping is performed when necessary. The pivot selection, row swapping decision, and elimination step control are handled by the coprocessor. Row swapping is implemented in local memory by moving and overwriting entire row vectors.

[0036] 3. Parallel Mapping and Multi-NPU Collaboration: To adapt to the vector parallelism width and on-chip storage organization of NPUs, the augmented matrix is ​​divided into several column-wise sub-blocks (e.g., N×16) and mapped to the local memory of multiple NPUs. Some NPUs store the column sub-blocks of the original matrix M, while the remaining NPUs store the column sub-blocks of the identity matrix I. Within the same elimination step, each NPU, based on the calculation instructions provided by the TPU operator library such as "tensor multiplication by constant" and "tensor subtraction," performs vectorized operations such as "row normalization" and "row elimination update" in parallel on different column sub-blocks of the same row. Since the column sub-blocks are distributed in parallel across multiple NPUs, different 16-dimensional sub-vectors of the same row can be scaled and subtracted synchronously on multiple NPUs, thereby updating all elements of the entire row in parallel. To ensure that each column sub-block has consistent addressing and data access rules during row transformations, the shape and stride of the tensor are set so that each column sub-block of the same row can participate in row operations according to a unified address calculation method.

[0037] 4. Operator reuse method: This parallel Gaussian elimination inversion operator can be used in the iterative process... Inverse matrix calculation can also be used for solving problems in the post-processing stage. The linear equations of a matrix (solved by transforming the linear equations into an augmented form or by reusing the elimination process) cover the key linear algebraic steps in the solution process within the same elimination operator framework.

[0038] (v) Local memory storage layout and workspace reuse 1. Unified step size alignment and layout convention: The matrix is ​​organized and stored in local memory using a unified aligned row step size, so that the matrix starting address and row step size meet the alignment constraint, thereby reducing the overhead of unaligned memory access. It also enables general matrix multiplication and Gaussian elimination inversion operators to directly read and write local memory data under the same data layout convention, reducing format conversion and data rearrangement.

[0039] 2. Partitioning and Reuse Strategy: Combining the local memory's bank-based organization method, the iterative matrix is... Temporary buffers, inverted augmented matrices, projection blocks, etc., are partitioned and planned. To ensure access bandwidth and the stability of parallel scheduling, a fixed local memory block is allocated for each type of key matrix for storage and computation. The same set of banks is time-division multiplexed at different stages: for example, in the iteration stage, it is used for storage... The temporary results are reused during the inversion phase to store the augmented matrix and intermediate elimination values. The switching timing of the above reuse is uniformly managed by the coprocessor, and through explicit lifecycle and read / write boundary control, it is ensured that data does not conflict between different steps.

[0040] (vi) Output and Interface This invention provides a unified calling interface and optional statistical information. Inputs include the system matrix (A, B, Q, R) and parameters (e.g., maximum number of iterations maxIter, convergence threshold). The output includes stable solutions. And the optimal feedback gain matrix K. Optional output information includes the number of iterations, residual index, and timing information, where the timing information can be used to calculate the overall operator scheduling time and NPU computation time separately for integration and evaluation.

[0041] (vii) Multi-system batch processing parallel strategy For solving the Riccati equations of multiple independent systems, a fully block-parallel strategy of "system-dimensional parallelism" is adopted: the host side processes the equations of N systems. The system is batch-packaged according to the system index and delivered to the device system memory at once. It also provides the number of systems, the parameters of each system dimension, and the base address or offset information of each system input / output, so that the device side can locate the data by index.

[0042] On the device side, the coprocessor performs unified system-level scheduling, mapping system tasks to multiple NPUs for parallel execution, specifically using one of the following two methods: 1. Parallel NPU allocation: Different systems are assigned to different NPUs (or NPU groups), enabling multiple systems to perform their core calculations such as matrix multiplication and inversion in parallel at the same time; 2. Batch Round-Robin Parallelism: When the number of systems exceeds the number of NPUs (or NPU groups) that can be parallelized, the system is divided into several batches. Each batch maps the system that does not exceed the available parallelism to the NPU for parallel solving. After the batch ends, the next batch of systems is processed in a round-robin fashion, thereby maintaining a stable system-level parallel execution structure under fixed hardware resources.

[0043] The solution process within each system reuses the end-to-end solution process of the single system, and the output of each system is... Write back to the Device system memory in system index order, and reclaim the memory on the Host side in the same order all at once. Complete batch processing output for multiple systems.

[0044] As a specific embodiment of this solution, the following is an example: Example 1 This embodiment demonstrates solving the continuous-time algebraic Riccati equations on a TPU platform, including: 1. Input and Output 1) Input matrix: System matrix , , , .

[0045] 2) Output: Riccati stable solution Optimal feedback gain .

[0046] The input / output data type is FP32, but FP16 / BF16, etc., can also be selected, depending on the device operator support.

[0047] 2. Host-side steps (input delivery and result collection) 1) Expand the input matrices A, B, Q, and R into a contiguous one-dimensional array in row-major order, and record the matrix order, shape / stride, and other metadata.

[0048] 2) Request and initialize the TPU device handle through the runtime library, and request system memory space on the device side.

[0049] 3) Write A, B, Q, and R into the device-side system memory at once, and construct the solution parameters; the parameters include at least: Input / output base address; Matrix dimensions n, m; Data type; Tensor shape / stride (including alignment step size settings); Maximum number of iterations: maxIter; Convergence threshold τ, etc.

[0050] 4) Start Riccati to solve the Device kernel, and complete the end-to-end solving process from the Device side.

[0051] 5) After the calculation is completed, the outputs P and K are read back from the system memory on the Device side to the Host side.

[0052] 3. Device-side steps (A53-Lite: Flow Control and Operator Scheduling) 1) The input matrix is ​​moved from system memory to on-chip local memory via GDMA, and the necessary data alignment / arrangement preparation is completed.

[0053] 2) On-chip local memory partitioning planning: A, B, Q, R, and iterative variables. On-chip space is allocated for inversion work area (augmented matrix buffer), projection block buffer, etc.; each matrix is ​​organized and stored using a uniform alignment row stride (e.g., aligned to 16), and partitioned and reused by bank / buffer.

[0054] 3) Construct the Hamiltonian matrix H and perform scaling preprocessing: Calculate the scaling factor α to obtain the initial value. α. The matrix multiplications involved are performed by the general GEMM operator on the NPU side.

[0055] 4) Iterative loop k=0,1,…,maxIter; a) Call the parallel Gaussian elimination inverse operator to calculate (or the solution to an equivalent linear equation); b) Obtained through iterative updates ; c) Calculate the amount required for the convergence criterion (e.g., relative change or residual), and use A53-Lite to determine whether the stopping criterion is met; if it is met, exit the iteration.

[0056] 5) After iterative convergence, construct... And based on this, a projection matrix is ​​formed. and its blocks (e.g.) , ).

[0057] 6) A stable solution P is recovered by solving a system of linear equations, and the solution can reuse the parallel elimination / inversion operator framework.

[0058] 7) Calculate the optimal feedback gain K (e.g., by a matrix multiplication chain; when the inverse or equivalent solution of R is involved, a parallel inversion operator can be called), and write P and K back to system memory.

[0059] 4. Device-side NPU sub-procedure: Parallel Gaussian elimination inversion operator (key step) 1) Constructing an augmented matrix in on-chip local memory ,in For the matrix to be inverted (e.g.) (or coefficient blocks in post-processing), where I is an identity matrix of the same order.

[0060] 2) Based on the NPU vector width and on-chip storage organization, the augmented matrix is ​​divided into multiple column fragments along the column direction (e.g., 16 columns per fragment; the fragment width is configurable), and mapped to the local memory of multiple NPUs: some NPUs store the column fragments of the left half M, and the remaining NPUs store the column fragments of the right half I.

[0061] 3) For each elimination step k: a) Each NPU reads the master information in parallel, and the A53-Lite determines the master and triggers row swapping when necessary; row swapping is implemented in local memory by moving entire rows of data. b) Normalize the pivot row (implemented by the NPU-side vectorized multiplication / scaling operators).

[0062] c) Perform row updates on non-pivotal rows (which can be achieved by a combination of operators such as "vectorized multiplication by constant" and "vectorized subtraction"). Since column slices are distributed across multiple NPUs, different column slices of the same row can be updated in parallel on multiple NPUs, thereby achieving parallel elimination of the entire row.

[0063] 4) After elimination, take the right half of the augmented matrix as... GDMA data interaction only occurs during the initialization loading and final write-back phases, and the elimination process is completed on-chip.

[0064] As another specific embodiment of this solution, the following is an example: Example 2 This embodiment demonstrates the implementation of fully block-parallel Riccati solving across multiple systems on a TPU platform. While maintaining the single-system end-to-end solving process of Embodiment 1, it adopts a fully block-parallel strategy of "system-dimensional parallelism" for solving scenarios involving multiple independent systems. This leverages the multi-NPU computing resources of the TPU to allocate different systems to different NPUs (or different NPU groups) for parallel execution, thereby improving batch processing throughput. This includes: 1. Batch Input Organization Method The host side will have N systems The system is batch-packaged according to the system index and delivered to the device system memory at once; at the same time, the number of systems, the parameters of each system dimension, and the base address or offset information of each system input / output are provided to facilitate the device side to locate the data by index.

[0065] 2. Multi-system parallel scheduling strategy On the device side, A53-Lite performs unified system-level scheduling, mapping system tasks to multiple NPUs for parallel execution. Specifically, one of the following methods can be used: 1) Parallel NPU allocation: Different systems are assigned to different NPUs (or NPU groups), enabling multiple systems to perform their core calculations such as matrix multiplication and inversion in parallel at the same time; 2) Batch Round-Robin Parallelism: When the number of systems exceeds the number of NPUs (or NPU groups) that can be parallelized, the system is divided into several batches. Each batch maps the system that does not exceed the available parallelism to the NPU for parallel solving. After the batch ends, the next batch of systems is processed in a round-robin fashion, thereby maintaining a stable system-level parallel execution structure under fixed hardware resources.

[0066] In the above parallel scheduling, the solution process within each system is reused from Example 1.

[0067] 3. Output organization and recycling Each system Write back to the Device system memory in system index order, and reclaim the memory on the Host side in the same order all at once. Complete batch processing output for multiple systems.

[0068] The above description is merely illustrative of the embodiments of the present invention and is not intended to limit the present invention. For those skilled in the art, any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A method for solving the optimal control Riccati equation based on TPU heterogeneous parallelism, characterized in that, Includes the following steps: S1: The host side organizes the system matrices A, B, Q, and R into continuous data in row-major order, writes them all at once into the device-side system memory of the tensor processor, and starts the Riccati equation solving kernel. After the calculation is completed, it reads back and outputs the stable solution matrix. With the optimal feedback gain matrix ; S2: On the device side, the coprocessor executes iterative process control and operator scheduling, moves the input matrix from the system memory to the on-chip local memory through global direct memory access, and completes data alignment and layout preparation; S3: Construct the Hamiltonian matrix in the on-chip local memory based on the input matrices A, B, Q, and R. Using the Frobenius norm pair Perform scaling preprocessing to obtain initial values ​​for iteration. ,in ;in, The Hamiltonian matrix is ​​constructed based on the system matrix. Describe the Frobenius norm. for The scaling factor is calculated using the Frobenius norm. This is the initial matrix for the symbolic function iteration; S4: Perform Newton-type iterations based on matrix sign functions. In each iteration, call the parallel Gaussian elimination inverse operator to calculate... The inverse matrix, according to the iterative update rule get The coprocessor calculates and determines whether the convergence criterion is met; if it is, the iteration exits. For the first The iteration matrix of the nth iteration. For the first The iteration matrix of the nth iteration. for The inverse matrix, The number of iterations ( ≥0); S5: After the iteration converges, construct the iterative result of the symbolic function. And based on this, a projection matrix is ​​formed. and its blocks , Solving linear equations by calling parallel Gaussian elimination inversion operator and general matrix multiplication operator. To restore the stable solution matrix ;in, Hamiltonian matrix The approximate value of the sign function, Represents matrix symbolic functions, For the projection matrix, To and Identity matrices of the same order This is the top-left block matrix of the projection matrix. This is the lower left block matrix of the projection matrix; S6: Based on optimal control relationship The optimal feedback gain matrix is ​​calculated using a general matrix multiplication operator and by calling a parallel Gaussian elimination inversion operator when needed. ,Will , Write back to the device-side system memory; among which, Input weight matrix The inverse matrix, For the input matrix The transpose of .

2. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to claim 1, characterized in that, In step S2, matrices A, B, Q, R, and the iteration matrix are stored in the on-chip local memory. The augmented matrix buffer and projection block buffer of the inversion work area are partitioned and planned. Each matrix is ​​organized and stored using a uniform alignment row stride, and partitioned and reused according to the storage body / buffer.

3. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to claim 1, characterized in that, In step S3, the Hamiltonian matrix H is constructed; in step S4, iterative updates are performed. Step S5: Block Calculation of Projection Matrix and Stable Solution Recovery, optimal feedback gain in step S6 All matrix multiplications in the computation process are uniformly executed by the general matrix multiplication operator on the neural network processing unit side of the tensor processor.

4. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to claim 1, characterized in that, The implementation process of the parallel Gaussian elimination inversion operator described in steps S4 and S5 is as follows: the matrix to be inverted is stored in on-chip local memory. Constructing augmented matrices ,in To and For identity matrices of the same order, Gaussian elimination transforms the left half of the augmented matrix into an identity matrix, and the right half becomes the inverse of the original matrix M. .

5. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to claim 4, characterized in that, In the Gaussian elimination process During the elimination process, the coprocessor selects a partial pivot selection strategy, choosing the element with the larger absolute value as the pivot, and performs row swapping when necessary. Row swapping is implemented in local memory by moving and overwriting data in whole row vectors. The pivot selection, row swapping decision and elimination step control are all handled by the coprocessor.

6. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to claim 5, characterized in that, augmented matrix The matrix is ​​divided into several column-wise segments, with the segment width adapted to the vector parallelism width of the neural network processing units. These column-wise segments are mapped to the local memory of multiple neural network processing units, with some units storing the original matrix. The columns are partitioned, and the remaining neural network processing units are stored in unit arrays. The column slices are divided into segments. Within the same elimination step, each neural network processing unit performs vectorized operations of row normalization and row elimination update in parallel on different column slices of the same row based on the tensor multiplication constant and tensor subtraction instructions provided by the tensor processor operator library. By setting the shape and step size of the tensor, each column slice participates in the row operation according to the unified address calculation rules.

7. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to claim 1, characterized in that, When solving the Riccati equations for multiple independent systems, the host side will have the matrices of N systems. Packaging is done in batches according to system index and distributed to the device-side system memory in one go; the device side allocates or reuses on-chip local memory partitions for multiple systems, and maps multiple systems to multiple neural network processing units or groups of neural network processing units to perform symbolic function iteration, projection block construction, and... , Recovery; when the number of systems exceeds the available parallel resources, a batch rotation method is used to perform fully block-based parallel solution of multiple systems.

8. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to claim 1, characterized in that, Input matrices A, B, Q, R and output matrix , The data type is single-precision floating-point, half-precision floating-point, or double-half-precision floating-point, depending on the operator support of the tensor processor device.

9. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to claim 1, characterized in that, The convergence criterion mentioned in step S4 is calculated as follows: ,in The preset convergence threshold, It is the Frobenius norm.

10. The optimal control Riccati equation solution method based on TPU heterogeneous parallelism according to any one of claims 2, characterized in that, On-chip local memory uses time-division multiplexing for its working area, and during the iteration phase, it is used to store the iteration matrix. The storage area for the augmented matrix and its temporary results is reused during the inversion phase to store the augmented matrix and intermediate quantities of elimination. The reuse switching timing is uniformly managed by the coprocessor, and data conflicts are avoided by clearly defining the data lifecycle and reading / writing boundary control.

Images