JAX分片数组上的离散差分计算：性能考量与实践

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 # 用于更精确的计时

def calc_fd_kernel(x):
    # 沿第一个轴计算一阶有限差分
    # prepend 参数用于在指定轴的前面添加值，以处理边界条件
    return jnp.diff(
        x, 1, axis=0, prepend=jnp.zeros((1, *x.shape[1:]))
    )

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数组进行分区
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()
}

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