加速大量2×2小型矩阵乘法运算的优化方案咨询

一、通过OpenBLAS配置多核心加速 #

1. 确认Numpy后端 #

先运行np.show_config()检查当前Numpy是否使用OpenBLAS。如果输出中无OpenBLAS相关信息，可通过conda安装OpenBLAS版本的Numpy：

conda install numpy openblas

2. 设置OpenBLAS线程数 #

OpenBLAS的线程数决定使用的CPU核心数，必须在导入numpy前设置环境变量：

import os
# 替换为你的CPU核心数，比如8
os.environ['OPENBLAS_NUM_THREADS'] = '8'
import numpy as np
import time

也可在运行脚本前通过系统环境变量设置：

• Linux/macOS：export OPENBLAS_NUM_THREADS=8 && python your_script.py
• Windows：set OPENBLAS_NUM_THREADS=8 && python your_script.py

二、优化多维数组运算逻辑（手动展开2×2矩阵乘法） #

原生@运算符对2×2小矩阵的乘法存在BLAS调用开销，手动展开矩阵乘法为元素级广播运算，能大幅提升效率。具体实现如下：

import os
os.environ['OPENBLAS_NUM_THREADS'] = '8'
import numpy as np
import time

A_dim = 6*6
B_dim = 2**8
C_dim = B_dim
A = np.random.rand(A_dim,A_dim,2,2)
B = np.random.rand(B_dim,2,2)
C = np.random.rand(C_dim,2,2)

tic = time.perf_counter()

# 提取所有矩阵的元素
# A的元素：(A_dim, A_dim)
A00, A01 = A[:, :, 0, 0], A[:, :, 0, 1]
A10, A11 = A[:, :, 1, 0], A[:, :, 1, 1]
# B的元素：(B_dim,)
B00, B01 = B[:, 0, 0], B[:, 0, 1]
B10, B11 = B[:, 1, 0], B[:, 1, 1]
# C的元素：(C_dim,)
C00, C01 = C[:, 0, 0], C[:, 0, 1]
C10, C11 = C[:, 1, 0], C[:, 1, 1]

# 计算 A@B 的元素：(B_dim, A_dim, A_dim)
AB00 = A00[None, :, :] * B00[:, None, None] + A01[None, :, :] * B10[:, None, None]
AB01 = A00[None, :, :] * B01[:, None, None] + A01[None, :, :] * B11[:, None, None]
AB10 = A10[None, :, :] * B00[:, None, None] + A11[None, :, :] * B10[:, None, None]
AB11 = A10[None, :, :] * B01[:, None, None] + A11[None, :, :] * B11[:, None, None]

# 计算 (A@B)@A 的元素：(B_dim, A_dim, A_dim)
ABA00 = AB00 * A00[None, :, :] + AB01 * A10[None, :, :]
ABA01 = AB00 * A01[None, :, :] + AB01 * A11[None, :, :]
ABA10 = AB10 * A00[None, :, :] + AB11 * A10[None, :, :]
ABA11 = AB10 * A01[None, :, :] + AB11 * A11[None, :, :]

# 计算 ((A@B)@A)@C 的元素：(B_dim, C_dim, A_dim, A_dim)
X00 = ABA00[:, None, :, :] * C00[None, :, None, None] + ABA01[:, None, :, :] * C10[None, :, None, None]
X01 = ABA00[:, None, :, :] * C01[None, :, None, None] + ABA01[:, None, :, :] * C11[None, :, None, None]
X10 = ABA10[:, None, :, :] * C00[None, :, None, None] + ABA11[:, None, :, :] * C10[None, :, None, None]
X11 = ABA10[:, None, :, :] * C01[None, :, None, None] + ABA11[:, None, :, :] * C11[None, :, None, None]

# 合并为最终的 (B_dim, C_dim, A_dim, A_dim, 2, 2) 张量
X = np.stack([np.stack([X00, X01], axis=-1),
              np.stack([X10, X11], axis=-1)], axis=-2)

toc = time.perf_counter()
print(f"matrix multiplication took {toc - tic:0.4f} seconds")

优化原理 #

2×2矩阵乘法的运算规则固定，手动展开后所有操作均为元素级广播运算，避免了BLAS对小矩阵乘法的调用开销，同时Numpy的元素运算能更高效地利用CPU向量化指令和多核心资源。

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