qmatmul：四元数矩阵快速乘法算法及其在串行与CUDA计算中的实现

《SoftwareX》：qmatmul: Fast multiplication of quaternion-valued matrices—algorithm and its implementations for sequential and CUDA computations

　　

研究人员提出了一种用于四元数（超复数，由一个实部和三个虚部构成）矩阵快速乘法的算法。相较于基于定义的公式，该算法所涉及的基本浮点乘法次数减半，且不受输入矩阵内容的影响。这一约减得益于合适的表示方式与分解策略，将目标乘积分解为两个子乘积，其一利用特定的对角对称性质，其二利用稀疏性。在软件相关部分，研究人员提出了八种实现变体（位于qmatmul Python模块中），并在约10⁶至10¹⁰量级的问题规模上进行了实验比较。这些变体涵盖：基于CPU的单线程与并行计算，以及基于GPU/CUDA的多线程计算。针对所提算法的CUDA计算设计包括：六个内核函数共十一次调用、合理的分块（tiling）与共享内存使用，以及少量的主机-设备内存传输。对于每种后端，研究人员报告了直接实现与所提算法实现之间的执行时间比值。对于足够大的问题规模，该比值趋近于浮点乘法理论上的两倍约减，而直接实现在小规模问题上仍具竞争力。此外，研究人员将所研究方法的基于cuBLAS的CuPy对应实现与基于Numba的自定义CUDA内核进行了基准测试。CuPy实现在float32计算中被证明极为高效，但在float64下往往被基于Numba的内核超越，尤其是对于偏斜矩形输入矩阵。

四元数（quaternion）是由一个实部和三个虚部构成的超复数，在三维几何与旋转计算中具有广泛应用，并延伸至计算机图形学、虚拟现实、机器视觉、机器人学、导航与控制理论、信号处理以及图论与组合数学等领域。然而，当代计算机科学与编程语言标准中并未提供四元数的内置数据类型，现有成熟软件库（如Eigen、Boost）虽支持基本四元数三维旋转与图形任务，却缺乏面向大规模四元数线性代数、特别是稠密四元数矩阵-矩阵乘法的优化实现。高性能超复数计算研究仍在进行中。既有研究如Williams-Young和Li从硬件中心视角研究了四元数高性能GEMM（通用矩阵-矩阵乘法）的实现条件，涉及指令缓存与存储缓存利用、SIMD（单指令多数据）与FMA（融合乘加）等微架构特性、多核CPU并行以及四元数表示方式比较，并形成了C++库HAXX，但未考虑在纯算法层面减少相对于定义式计算的运算次数。Cariow提出了利用矩阵-向量或矩阵-矩阵乘积中块对称性的通用代数框架，Cariow、Naumowicz和Handkiewicz则针对单个四元数乘积开展了面向硬件的研究（以FPGA模块实现并综合），但均未涉及矩阵-矩阵乘积的一般情形。基于此，研究人员开展了四元数矩阵快速乘法的研究。

研究人员的主要贡献包括：(i) 适用于书写一般矩阵-矩阵乘积的四元数矩阵堆叠表示（stacked representation）；(ii) 将矩阵-矩阵乘积分解为两个子乘积，记为A?·B⁴与A?·B⁴，为更快的计算奠定基础；(iii) 利用A?的双重嵌套对角对称性与A?的稀疏性；(iv) 基于上述要点并结合合适的Hadamard变换与块置换，形成主快速算法（算法1）；(v) 算法的NumPy/BLAS实现；(vi) 算法的多线程CUDA计算设计与实现（6个内核函数组织为11次启动）；(vii) 在三种GPU设备上、约10⁶至10¹⁰量级问题规模下对多种计算后端的充分效率测试；(viii) 补充内容包括乘法/加法/FLOPs的精确计数、内存占用分析、QCNN应用示例、向CuPy（cuBLAS）的转换及效率比较；(ix) 数值精度分析。研究得出结论：该算法的关键性质是将基本浮点乘法次数相对于定义公式减半，计算实验证实了这一点；直接实现与基于算法实现的比值随问题规模增大而趋近理论上的两倍约减，而在小规模下因实现与调度开销占主导而接近或低于1；所设计的多线程CUDA实现在大输入矩阵下最为高效。qmatmul软件包以Python实现，输入四元数矩阵以NumPy三维数组表示（最后一维大小为4，存储每个四元数的实部与虚部），便于融入Python生态系统，且借助Numba/LLVM即时编译，潜在应用无需集成原生CUDA C++ API即可完全嵌入Python生态。该研究为四元数相关应用（三维旋转、几何、物理、机器人学、神经网络等）提供了高效计算工具，并指出QCNN（四元数卷积神经网络）研究是可能的未来方向。

研究人员开展研究用到的主要关键技术方法包括：基于堆叠表示与代数分解的快速矩阵乘法算法设计，利用Hadamard变换（H₄?I）与块置换实现对角对称性与稀疏性的利用；基于NumPy/BLAS的CPU单线程与多线程实现（通过threadpoolctl限制或放开BLAS线程数）；基于Numba即时编译（@cuda.jit → IR → LLVM/NVVM → PTX → SASS）的CUDA多线程内核实现（采用tiling与共享内存、tile_size = 8）；以及基于CuPy（cuBLAS后端）的对照实现。实验在三种计算环境（分别配备GeForce RTX 3090、GeForce RTX 5090、RTX PRO 5000 Blackwell GPU）上进行，问题规模由M·N·P定义，从1.0·10⁶至2.7·10¹⁰，输入矩阵由qmatmul.py模块中的qmatrand辅助函数生成（取值范围(?2.0, 2.0)，随机种子为0），正式计时前以M=N=P=100的小问题进行一次预热调用以触发Numba即时编译。

研究结果方面：（1）效率测试结果。研究人员对八种计算途径进行了实验，包括基于定义式的Numba单线程与并行实现、基于直接公式的NumPy单/多线程实现、直接公式的CUDA实现，以及快速算法的NumPy单/多线程与CUDA实现。结果表明，基于定义式的单线程基线方法在所有实验中最为缓慢；NumPy-based方法轻易优于朴素方法，即便单线程NumPy直接实现也快于Numba并行朴素实现，这归因于高度优化的BLAS计算；对于较大问题规模，qmatmul_algo_numpy_st相对qmatmul_direct_numpy_st的执行时间约缩短1.8倍，与理论预期几乎完全吻合；NumPy并行变体性能难以明确评估，因BLAS线程池管理开销而呈现较高方差，且在大规模下算法相对直接实现的优势降至约1.2倍；CUDA实现中，直接法_direct_numba_cuda在小规模下即具竞争力，而算法法_algo_numba_cuda在问题规模?10⁹时明确且一致地最优，相对基线的最高加速在环境2中分别约为float32的15300倍与float64的6550倍；对于大规模问题，两种CUDA实现的执行时间比值趋近或超过2.0。研究人员解释，变换矩阵形式的直接法CUDA实现仅发出4个内核、共11次CUDA API调用，而所提算法使用11个内核、共27次API调用，调度开销相差2.45倍，这解释了小规模下直接实现因每次调用调度开销占主导而保持竞争力的现象；在设备内存方面，直接实现的总占用为20MN+8NP+8MP个元素，所提算法为16MN+16NP+20MP个元素，当M=N=P时所提算法总设备内存约多44%。（2）QCNN前向传播应用示例结果。研究人员给出一个小型QCNN前向传播的概念验证实现，网络含两个卷积层、展平与一个全连接层，作用于B=32的合成64×64 RGB-D图像批次，权重与特征图均为四元数；每个卷积以im2col式展开后接单次qmm.dot调用实现，权重采用四元数适配的He-uniform或Glorot-uniform初始化，卷积后应用分量式QReLU。逐层计时显示与前述相同的交叉模式：conv1处于小N区间，直接法仍具竞争力（比值≈1.0）；conv2（3×3核）直接法已比算法法多耗时31%（比值1.31）；全连接层（N=262144）差距最大，直接法比算法法多耗时88%（比值1.88），趋近算法的渐近行为；对全连接层输出与直接法进行正确性检查，相对Frobenius误差为1.58·10^?5，最大绝对误差为5.71×10^?5（输出典型尺度RMS为8.15×10^?1），numpy.allclose(rtol=1e-5, atol=1e-2)为True。（3）CuPy实现与比较结果。研究人员将_direct_numba_cuda与_algo_numba_cuda直译为CuPy实现，其中直接法所依据的公式(15)实为单次实值GEMM，而算法1由八次实值GEMM组成（每个对角乘积D(·,·)四次）。在环境3中对四类问题形状与两种数据类型的计时比较表明，CuPy实现在float32下对所有四类形状均优于两种Numba-CUDA实现，对形状类别1，algo_cupy比基线快9.8倍、比algo_numba_cuda快4.2倍；而在float64下，对形状类别2–4，CuPy反而不及基于Numba的实现，对QCNN全连接层对应的形状类别2，algo_cupy比direct_numba_cuda慢2.9倍、比algo_numba_cuda慢5.1倍，最后两个类别证实了这一趋势，最极端的类别4表明algo_numba_cuda的结果更为有利。

讨论部分总结了上述结果：基线定义式方法最慢，其CPU并行变体快约10至70倍；NumPy-based方法因BLAS高度优化而表现优异，单线程直接与算法实现的头对头比较在大规模下与理论预期近乎完美吻合，而NumPy并行变体因线程池管理开销存在高方差且算法优势被削弱，原因在于算法并行实现涉及两个较小的dot运算与多达十个辅助任务，造成工作负载碎片化与频繁同步屏障；CUDA实现中两种设计形成对比，直接法内核少、API调用少，算法法内核与API调用多但在大规模下计算优势显现，且设备内存占用更高；此外，附录B针对算法因Hadamard类变换、置换、缩放与最终减法等算术重组可能影响浮点数值精度的问题，将float32计算与float64参考进行了比较。研究结论部分指出：研究人员考虑了可能规模显著较大的四元数矩阵-矩阵乘法任务并提出了一种快速算法，该算法利用输入矩阵的堆叠表示、将目标乘积分解为两个分别呈现对角对称性与稀疏性的合适子乘积、Hadamard类变换与块置换等思想，其关键性质是将基本浮点乘法次数相对于定义公式减半，这一性质得到了计算实验的证实；直接实现与基于算法实现的比值随问题规模增大而证实理论上的两倍约减，而在小规模下因实现与调度开销占主导而保持接近或低于1；除自然的NumPy/BLAS实现外，研究的相当部分致力于为其设计多线程CUDA实现，该实现在大输入矩阵下被证明最为高效。研究人员认为该项目具有较高潜在影响力：qmatmul软件包以Python实现，输入四元数矩阵以NumPy三维数组表示，便于融入Python生态系统；使用简便，矩阵乘法归结为对qmm.dot的简单调用；效率由所提算法本身与GPU多线程CUDA计算共同保证；且借助Numba/LLVM即时编译，潜在应用无需集成原生CUDA C++ API即可完全嵌入Python生态。研究人员表示希望看到其他研究者或工程师利用qmatmul开展工作，并将QCNN研究（在总参数数量相近时QCNN能否在某些应用中优于传统CNN、四元数矩阵-矩阵乘积能否成功支撑QCNN随机梯度下降中的前向与反向计算）列为可能的未来研究方向。