JAX分片数组上的离散差分计算：性能考量与优化策略

import os
os.environ["XLA_FLAGS"] = (
    f'--xla_force_host_platform_device_count=8'
)

import jax as jx
import jax.numpy as jnp
import jax.experimental.mesh_utils as jxm
import jax.sharding as jsh
import timeit # For performance measurement

def calc_fd_kernel(x):
    # Calculate 1st-order fd along the first axis
    # prepend with zeros to match output shape, maintaining original shape logic
    return jnp.diff(
        x, 1, axis=0, prepend=jnp.zeros((1, *x.shape[1:]), dtype=x.dtype)
    )

def make_fd(shape, shardings):
    # Compiled fd kernel factory
    return jx.jit(
        calc_fd_kernel,
        in_shardings=shardings,
        out_shardings=shardings,
    ).lower(
        jx.ShapeDtypeStruct(shape, jnp.dtype('f8')) # Define input shape and dtype for AOT compilation
    ).compile()

# Create 2D array to partition
n = 2**12 # e.g., 4096
shape = (n, n,)

x = jx.random.normal(jx.random.PRNGKey(0), shape, dtype='f8')

# Define device mesh and sharding strategies
# Use all available CPU devices
devices = jx.devices("cpu")
if len(devices) < 8:
    print(f"Warning: Only {len(devices)} CPU devices available. Some sharding configurations might not be fully utilized.")
    # Adjust for available devices if less than 8
    num_devices_to_use = min(8, len(devices))
else:
    num_devices_to_use = 8

shardings_test = {
    (1, 1) : jsh.PositionalSharding(jxm.create_device_mesh((1,), devices=devices[:1])).reshape(1, 1),
    (num_devices_to_use, 1) : jsh.PositionalSharding(jxm.create_device_mesh((num_devices_to_use,), devices=devices[:num_devices_to_use])).reshape(num_devices_to_use, 1),
    (1, num_devices_to_use) : jsh.PositionalSharding(jxm.create_device_mesh((num_devices_to_use,), devices=devices[:num_devices_to_use])).reshape(1, num_devices_to_use),
}

# Place arrays onto devices according to sharding
x_test = {
    mesh_config : jx.device_put(x, shardings)
    for mesh_config, shardings in shardings_test.items()
}

# Compile the fd kernel for each sharding strategy
calc_fd_test = {
    mesh_config : make_fd(shape, shardings)
    for mesh_config, shardings in shardings_test.items()
}

# Measure execution time for each configuration
print("Measuring performance for different sharding strategies:")
for mesh_config, x_sharded in x_test.items():
    calc_fd_compiled = calc_fd_test[mesh_config]
    print(f"\nConfiguration: {mesh_config}")
    # Use a lambda to ensure the function is called with the specific sharded array
    # and block_until_ready() to wait for all computations to complete
    stmt = f"calc_fd_compiled(x_sharded).block_until_ready()"
    # Use globals for timeit to access calc_fd_compiled and x_sharded
    globals_dict = {"calc_fd_compiled": calc_fd_compiled, "x_sharded": x_sharded}
    # timeit.repeat to get multiple runs for better statistics
    times = timeit.repeat(stmt, globals=globals_dict, number=1, repeat=7)
    print(f"{min(times)*1000:.3f} ms ± {jnp.std(jnp.array(times))*1000:.3f} ms per loop (min ± std. dev. of 7 runs, 1 loop each)")

Configuration: (1, 1)
48.9 ms ± 414 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Configuration: (8, 1)
977 ms ± 34.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Configuration: (1, 8)
48.3 ms ± 1.03 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)