Matrix Power

原始题目：LeetGPU - Matrix Power

题目描述

编写一个 GPU 程序，将一个 $N \times N$ 的方阵 $A$ 提升到整数次幂 $P$。solve 函数接收扁平化的输入矩阵 input（行优先）、一个同样大小的空输出矩阵 output、维度 $N$ 和指数 $P$。计算：

$$output = A^P$$

其中矩阵乘法使用标准的稠密 32 位浮点乘法。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须以行优先顺序写入 output 数组。

示例

示例 1

Input:  A = [[1,2],[3,4]], N = 2, P = 3
Output: [[37,54],[81,118]]

示例 2

Input:  A = [[1,0,2],[0,1,0],[3,0,0]], N = 3, P = 2
Output: [[7,0,2],[0,1,0],[3,0,6]]

约束条件

• $1 \le N \le 1{,}024$，$1 \le P \le 20$。
• $-10.0 \le A_{ij} \le 10.0$。
• 性能测试在 $N = 512$ 的规模下进行。

解题思路

矩阵幂可以使用二分快速幂：$A^P = (A^2)^{P/2}$（$P$ 为偶数）或 $A \cdot A^{P-1}$（$P$ 为奇数），将复杂度从 $O(P)$ 降到 $O(\log P)$。每次乘法是一个标准的 GEMM，可以复用之前的分块共享内存策略。由于 $P \le 20$，最多只需约 5 次矩阵乘法，计算成本可控。欢迎在 GitHub Discussions 分享你的解法。