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 次矩阵乘法，计算成本可控。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void matmul(const float* A, const float* B, float* C, int N) {
    int r=blockIdx.y*blockDim.y+threadIdx.y, c=blockIdx.x*blockDim.x+threadIdx.x;
    if(r<N&&c<N){float s=0;for(int k=0;k<N;k++)s+=A[r*N+k]*B[k*N+c];C[r*N+c]=s;}
}
extern "C" void solve(const float* input, float* output, int N, int P) {
    // Copy input to output as identity iteration
    float *tmp1,*tmp2,*src,*dst;
    cudaMalloc(&tmp1,N*N*sizeof(float));cudaMalloc(&tmp2,N*N*sizeof(float));
    cudaMemcpy(tmp1,input,N*N*sizeof(float),cudaMemcpyDeviceToDevice);
    cudaMemset(output,0,N*N*sizeof(float));
    for(int i=0;i<N;i++)output[i*N+i]=1.0f; // Identity
    dim3 t(16,16),b((N+15)/16,(N+15)/16);
    for(int p=0;p<P;p++){matmul<<<b,t>>>(output,input,tmp1,N);cudaMemcpy(output,tmp1,N*N*sizeof(float),cudaMemcpyDeviceToDevice);}
    cudaFree(tmp1);cudaFree(tmp2);cudaDeviceSynchronize();
}

Triton

import torch
def solve(A, P):
    return torch.linalg.matrix_power(A, P)