如何并行化多组矩阵乘法？含GPU异步/并行实现问询

问题背景 #

现有如下串行执行的多组矩阵乘法代码：

import numpy as np

a = np.random.rand(10, 20)
b = np.random.rand(20, 1000)

ans_list = []

for i in range(980):
    ans_list.append(
        np.dot(a, b[:, i:i+20])
    )

已知NumPy会并行化单个矩阵乘法，但需要实现外层循环的并行，让多组乘法同时执行；同时需要了解使用CuPy（GPU替代NumPy）时，如何异步/并行提交多组矩阵乘法任务。

补充说明：上述滑动窗口仅为示例，已知该特例可通过np.lib.stride_tricks.sliding_window_view优化，但需求是任意多组矩阵乘法的最优并行方案，而非该特例优化。

CPU端并行方案 #

1. 多进程实现（multiprocessing.Pool） #

受Python GIL限制，CPU密集型任务适合用多进程突破单线程瓶颈：

import numpy as np
from multiprocessing import Pool

def single_dot(args):
    a, b_slice = args
    return np.dot(a, b_slice)

if __name__ == "__main__":
    a = np.random.rand(10, 20)
    b = np.random.rand(20, 1000)
    
    # 生成所有计算任务的参数对
    tasks = [(a, b[:, i:i+20]) for i in range(980)]
    
    # 按CPU核心数设置进程池大小
    with Pool(processes=4) as pool:
        ans_list = pool.map(single_dot, tasks)

• 注意：进程间数据拷贝存在开销，若矩阵体积过大，可使用multiprocessing.Array等共享内存机制优化。

2. 高层多进程API（ProcessPoolExecutor） #

更简洁的封装，用法与Pool类似：

import numpy as np
from concurrent.futures import ProcessPoolExecutor

def single_dot(a, b_slice):
    return np.dot(a, b_slice)

if __name__ == "__main__":
    a = np.random.rand(10, 20)
    b = np.random.rand(20, 1000)
    
    tasks = [(a, b[:, i:i+20]) for i in range(980)]
    
    with ProcessPoolExecutor(max_workers=4) as executor:
        ans_list = list(executor.map(single_dot, *zip(*tasks)))

GPU端（CuPy）并行/异步方案 #

CuPy天然适配GPU并行计算，同时支持异步执行提升任务吞吐量。

1. 异步流控制 #

通过创建多个CUDA流，让多组矩阵乘法并行排队执行：

import cupy as cp

a = cp.random.rand(10, 20)
b = cp.random.rand(20, 1000)

ans_list = []
streams = []

for i in range(980):
    # 创建独立CUDA流
    stream = cp.cuda.Stream()
    streams.append(stream)
    with stream:
        # 异步执行矩阵乘法，结果在流中排队
        res = cp.dot(a, b[:, i:i+20])
        ans_list.append(res)

# 等待所有流执行完成
for stream in streams:
    stream.synchronize()

# 可选：将结果转回CPU
ans_list_cpu = [cp.asnumpy(res) for res in ans_list]

2. 批量矩阵乘法（最优方案） #

将多组待计算矩阵整理为3D数组，一次性完成批量计算，GPU利用率远高于循环异步提交：

import cupy as cp

a = cp.random.rand(10, 20)
b = cp.random.rand(20, 1000)

# 生成所有切片并堆叠为(980, 20, 20)的3D数组
b_slices = cp.lib.stride_tricks.sliding_window_view(b, (20, 20)).squeeze()

# 批量矩阵乘法：(10,20) @ (980,20,20) -> (980,10,20)
ans_array = cp.tensordot(a, b_slices, axes=1)
ans_list = list(ans_array)

• 适配任意多组任务：只要能将多组右矩阵整理为3D数组，即可用该方式实现高效批量计算。

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