JAX分片机制概述
JAX是一个强大的高性能数值计算库，它利用即时编译（JIT）和自动微分等技术实现高效的科学计算。为了进一步提升大规模计算的性能，JAX引入了分片（Sharding）机制。分片允许用户将大型数组逻辑上分割成更小的块（shard），并将这些块分布到多个计算设备（如CPU核心、GPU或TPU）上。JAX内部会自动处理这些分片上的计算分布，旨在通过并行化来加速整体运算。
离散差分操作的挑战
离散差分（Discrete Difference）是一种常见的数值操作，例如计算一阶差分 f(x_i) - f(x_{i-1})。这类操作具有局部数据依赖性，即计算某个元素的值需要其相邻元素。当数组被分片并分布到不同设备上时，如果相邻元素恰好位于不同的设备上，那么计算就需要跨设备通信来获取所需数据。这种通信开销可能抵消甚至超过并行计算带来的好处。
为了探究JAX分片对离散差分计算性能的影响，我们设计了一个实验，在多核CPU环境下测试不同的分片策略。
实验设置与代码实现
本实验旨在通过JAX的自动并行化功能，评估在分片数组上执行一阶离散差分操作的性能。我们使用AOT（Ahead-Of-Time）编译JAX NumPy API调用，并测试了多种设备网格配置，以观察沿差分方向和垂直于差分方向的分片效果。
首先，设置XLA环境变量以强制JAX使用多个CPU设备。这里我们配置了8个CPU核心作为可用的计算设备。
import os
os.environ["XLA_FLAGS"] = (
f'--xla_force_host_platform_device_count=8' # 使用8个CPU核心作为设备
)
import jax as jx
import jax.numpy as jnp
import jax.experimental.mesh_utils as jxm
import jax.sharding as jsh
import time # 用于更精确的计时接下来，定义一个计算一阶离散差分的核函数。为了处理边界条件，我们使用 jnp.diff 并预先添加一行零，确保差分操作的输出形状与输入匹配。
def calc_fd_kernel(x):
# 沿第一个轴计算一阶有限差分
# prepend 参数用于在指定轴的前面添加值，以处理边界条件
return jnp.diff(
x, 1, axis=0, prepend=jnp.zeros((1, *x.shape[1:]))
)为了利用JAX的AOT编译和分片功能，我们创建一个工厂函数 make_fd。这个函数根据给定的数组形状和分片规则编译 calc_fd_kernel。in_shardings 和 out_shardings 参数是实现自动并行化的关键，它们告诉JAX如何将输入和输出数组分片到设备上。
def make_fd(shape, shardings):
# 编译差分核函数的工厂函数
# jx.jit 编译函数，in_shardings 和 out_shardings 指定输入输出的分片方式
return jx.jit(
calc_fd_kernel,
in_shardings=shardings,
out_shardings=shardings,
).lower(
jx.ShapeDtypeStruct(shape, jnp.dtype('f8')) # 定义输入数组的形状和数据类型
).compile() # 进行AOT编译然后，我们创建一个大型的2D数组作为实验数据，并定义三种不同的分片策略：无分片（基线）、沿差分轴分片和垂直于差分轴分片。
# 创建一个2D数组进行分区
n = 2**12 # 数组维度设置为 4096x4096
shape = (n,n,)
# 生成随机数据作为输入数组
x = jx.random.normal(jx.random.PRNGKey(0), shape, dtype='f8')
# 定义不同的分片测试配置
shardings_test = {
# (1, 1): 无分片，所有数据在一个设备上
# 创建一个1x1的设备网格，所有数据都在第一个CPU设备上
(1, 1,) : jsh.PositionalSharding(jxm.create_device_mesh((1,), devices=jx.devices("cpu")[:1])).reshape(1, 1),
# (8, 1): 沿第一个轴（差分方向）分片到8个设备
# 创建一个8x1的设备网格，将数据沿第一个轴分片到8个CPU设备
(8, 1,) : jsh.PositionalSharding(jxm.create_device_mesh((8,), devices=jx.devices("cpu")[:8])).reshape(8, 1),
# (1, 8): 沿第二个轴（垂直于差分方向）分片到8个设备
# 创建一个1x8的设备网格，将数据沿第二个轴分片到8个CPU设备
(1, 8,) : jsh.PositionalSharding(jxm.create_device_mesh((8,), devices=jx.devices("cpu")[:8])).reshape(1, 8),
}
# 将原始数组根据不同的分片规则放置到设备上
x_test = {
mesh : jx.device_put(x, shardings) # jx.device_put 将数据放置到指定分片规则的设备上
for mesh, shardings in shardings_test.items()
}
# 为每种分片配置编译差分函数
calc_fd_test = {
mesh : make_fd(shape, shardings)
for mesh, shardings in shardings_test.items()
}最后，我们对每种分片配置下的差分计算进行计时，以评估其性能。为了确保计时准确，我们使用 block_until_ready() 等待所有计算完成。
print("--- 性能测试结果 ---")
for mesh_config, x_sharded in x_test.items():
calc_fd_compiled = calc_fd_test[mesh_config]
start_time = time.perf_counter()
result = calc_fd_compiled(x_sharded).block_until_ready() # 确保计算完成
end_time = time.perf_counter()
print(f"分片配置 {mesh_config}: {(end_time - start_time) * 1000:.3f} ms")
# 原始问题中的%timeit结果示例，供参考：
# (1, 1) - 无分片
# 48.9 ms ± 414 µs per loop 
