向量化编程与广播（1）：引言

著作信息

• 原作者：陈久宁

• 最早于 2022-05-15 00:19 发布于 JuliaCN 公众号

前言

当你执行 a .* b 时，你就在使用向量化编程与广播。本文将以绘制二维曲面为例介绍广播这一概念。

例：绘制二维曲面

绘制 f(x,y) = x*exp(-x^2-y^2) 的曲面是一个典型的使用广播的例子。简单来说，为了绘制该函数，我们需要：

1. 将坐标 (x,y) 离散化为格点

2. 计算每个格点上的函数值 f(x,y)

因此，如果我们使用 for 循环来做这件事情的话，它大概会是下面这样子：

using Plots

f(x, y) = x*exp(-x^2-y^2)
X = -2.0:0.1:2.0
Y = -2.0:0.1:2.0

Z = zeros(length(X), length(Y))
for i in axes(X, 1), j in axes(Y, 1)
	Z[i,j] = f(X[i], Y[j])
end
surface(Z)

对于 MATLAB 用户来说，这件事情通常会使用 meshgrid 函数来实现：

% matlab
X = -2.0:0.1:2.0;
Y = -2.0:0.1:2.0;
 
[X, Y] = meshgrid(X, Y);
Z = X .* exp(-X.^2-Y.^2);
 
surf(X, Y, Z)

Python 用户也可以在 numpy 下找到 meshgrid 函数，似乎这就足够了。除此之外，告诉你 meshgrid 函数的人也许也会告诉你使用 meshgrid 的效率非常高。

那么问题来了：

1. 如何理解 meshgrid 以及为什么要使用 meshgrid？

2. Julia 里面存在 meshgrid 的等价功能吗？

为了说明这个问题，让我们用 Julia 实现一个简单的 meshgrid 函数:

function meshgrid(X, Y)
	X̂ = repeat(reshape(X, :, 1), 1, length(Y))
	Ŷ = repeat(reshape(Y, 1, :), length(X), 1)
	return X̂, Ŷ
end

这样做使得循环时矩阵的指标可以保持一致，即下面的两个版本在结果上是等价的：

1: 原始版本

for i in axes(X, 1), j in axes(Y, 1)
	Z[i, j] = f(X[i], Y[j])
end

2: meshgrid + 笛卡尔下标

X̂, Ŷ = meshgrid(X, Y)
for i in axes(X, 1), j in axes(X, 2)
	Z[i, j] = f(X̂[i, j], Ŷ[i, j])
end

更进一步地，如果我们使用线性下标的话，第二个版本又等价于：

3: meshgrid + 线性下标

X̂, Ŷ = meshgrid(X, Y)
for i in 1:length(X̂) # X̂ 的尺寸为 (length(X), length(Y))
	Z[i] = f(X̂[i], Ŷ[i])
end

聪明的小脑袋瓜们已经开始思考了：创建这样的一个辅助矩阵来统一下标，除了浪费内存空间以外，有什么实用价值吗？这是个好问题。实际上，meshgrid 是一个很古老的 MATLAB 函数：2006 年之前它就存在了。在那时它用来构造下标一致的矩阵，从而辅助向量化编程的进行来得到高性能的代码。然而随着技术发展到现在，meshgrid 已经成了一个不必要的功能，而且正如下面的性能测试所表明的，使用meshgrid还会降低代码速度。

在这里我们第一次遇到了向量化编程的核心手段：「广播 (broadcast)」。简单来说，在不引起歧义的情况下，广播允许不同尺寸的数组混合在一起进行逐点运算。这是什么意思呢？在 Julia 中，对于 meshgrid 这个特定例子下面两种写法基本等价：

4: 广播 -- 利用点运算符 . 来触发自动的内存分配和循环

f_vec(X, Y) = f.(X, Y)

function f_loop(X, Y)
	Z = similar(X, eltype(X), (length(X), length(Y)))
	@inbounds for i in 1:length(X) # @inbounds 用来关闭下标数组越界检查
		@simd for j in 1:length(Y) # @simd 用来触发 CPU SIMD 并行指令集
			Z[i,j] = f(X[i], Y[j])
		end
	end
	return Z
end

我们分别比较一下上面提到的等价版本在向量长度为 64 与 512 规模下的执行时间：

其中 LoopVectorization 用于优化 for 循环代码，关于它的使用以及背后的设计涉及到的东西非常多，我们也留到未来用多篇文章来单独介绍。

5: 利用 LoopVectorization 进一步优化 for 循环

using LoopVectorization

function f_vec!(Z, f, X, Y)
	@turbo for i in axes(X, 1)
		for j in axes(Y, 1)
			Z[i, j] = f(X[i], Y[j])
		end
	end
	return Z
end

从上面的性能对比我们可以得到这样的结论：

1. Julia 的 for 循环与广播的版本性能基本一致

2. meshgrid 在现在已经是一个不必要的累赘

3. MATLAB 的 for 循环非常缓慢，因此高性能代码必须使用广播 （向量化编程）

4. Julia 比 MATLAB 要慢（划掉） Julia 的广播实现目前仅仅只利用了 SIMD 指令而没有更进一步地利用 AVX 指令，因此比 MATLAB 更进一步优化后的广播实现要慢。并且我们可以期待未来 Julia 将 LoopVectorization 引入到广播的实现中，从而得到远超 MATLAB 的性能。

在下一篇文章，我们将更具体地介绍广播的基本规则。