如何用CuPy将矩阵伪逆应用于N个数组？GPU性能不及CPU原因解析

场景与测试环境 #

需要对形状为(100000, 1397, 2)的主数组应用伪逆运算，测试硬件配置：

• CPU：Intel 24核13900k
• GPU：Nvidia RTX 4090

数据生成代码：

import numpy as np
import cupy as cp

# generating the data
arr = np.random.uniform(low=0.5, high=1500.20, size=(1397, 2))
main_arr = np.tile(arr, (100000, 1, 1))

NumPy版本测试结果 #

直接调用np.linalg.pinv处理3D数组，耗时11秒：

%%time
np.linalg.pinv(main_arr)

执行结果：

CPU times: user 22.5 s, sys: 27.4 s, total: 49.9 s
Wall time: 11 s

CuPy版本问题与测试结果 #

直接调用cp.linalg.pinv处理3D数组报错，因为CuPy的pinv仅支持二维数组：

main_arr_gpu = cp.array(main_arr) # Copy the array to the GPU
cp.linalg.pinv(main_arr_gpu)

报错信息：

LinAlgError: 3-dimensional array given. Array must be two-dimensional

改用列表推导式逐个处理子数组，耗时22.3秒（未计入数据传输时间）：

%%time
[cp.linalg.pinv(arr_gpu) for arr_gpu in main_arr_gpu]

性能差异原因分析 #

1. NumPy的批量并行优化
   NumPy的linalg.pinv对高维数组做了原生批量处理优化，底层调用的BLAS库会充分利用CPU的多核并行能力（13900k的24核能高效分摊计算负载），且一次性处理所有子数组，无循环调度开销。

2. CuPy的循环调度开销
   CuPy的pinv不支持3D数组，只能逐个处理每个(1397,2)子数组。Python列表推导式本质是串行循环，每个cp.linalg.pinv调用都会触发一次GPU kernel启动，10万次启动的延迟累加后，完全抵消了GPU单任务的计算优势。

3. 任务粒度不匹配GPU特性
   GPU擅长处理大规模、高计算量的并行任务，而每个(1397,2)子数组的伪逆运算属于轻量计算任务，kernel启动延迟在总耗时中占比极高，反而不如CPU多核并行处理高效。

4. 内存访问效率差异
   NumPy数组在CPU内存中是连续布局，BLAS库能充分利用CPU缓存提升访问效率；CuPy数组在GPU显存中，轻量任务的内存访问模式无法充分利用高显存带宽，进一步放大了性能差距。

优化建议 #

使用CuPy的vmap（向量映射）功能批量处理子数组，将10万次kernel调用合并为一次，大幅降低调度延迟：

from cupyx.scipy.linalg import pinv

# 定义单个数组的伪逆运算函数
def single_pinv(arr):
    return pinv(arr)

# 用vmap实现批量处理，in_axes=0表示对第0维度的每个元素应用函数
vmap_pinv = cp.vmap(single_pinv, in_axes=0)
result = vmap_pinv(main_arr_gpu)

内容的提问来源于stack exchange，提问作者Begoodpy