Julia中Y=A时如何高效实现Y*=B矩阵乘法？

问题描述 #

矩阵乘法 Y = A * B 可以通过 mul!(Y, A, B) 实现零内存分配，但当需要计算 Y *= B（即 Y = Y * B）时，mul! 无法直接使用（输入和输出同为Y会覆盖中间计算结果）。请问是否存在类似高效的方法实现该操作？如果没有，最高效的实现方式是什么？

测试示例 #

n = 10
A = rand(n,n)
B = rand(n,n)
Y = zeros(n,n)

# mul! 无额外内存分配
@allocated Y = A * B   # 输出896
@allocated mul!(Y, A, B) # 输出0    

# mul! 无法处理Y作为输入和输出的情况
@allocated Y *= B # 输出896

# 期望的零分配函数效果
@allocated mul_2!(Y, B)  # 输出0

解决方案 #

Julia 的 LinearAlgebra 模块中直接提供了原地右乘函数 rmul!(Y, B)，完全满足零内存分配的需求，它就是你示例中设想的 mul_2!(Y, B) 的官方实现。

验证代码 #

using LinearAlgebra

n = 10
A = rand(n,n)
B = rand(n,n)
Y = copy(A)  # 初始化Y为A的副本

# 测试rmul!的零分配特性
@allocated rmul!(Y, B)  # 输出0

原理与最优实现说明 #

rmul!(Y, B) 底层调用了经过高度优化的 BLAS 原地乘法接口，直接在 Y 的内存空间上完成计算，不会额外分配新矩阵，性能和 mul! 完全一致。

如果需要手动实现，最优方式是基于 BLAS 的原地操作进行封装，但官方的 rmul! 已经是工业级的最优实现，无需重复造轮子。

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