guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-06-29 16:41:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use new cic_paint with halo

Browse source
This commit is contained in:
Wassim KABALAN 2024-07-18 12:42:07 +02:00
parent 5775a37550
commit 7501b5bc6d
2 changed files with 77 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

82
jaxpm/distributed.py
View file

@ -10,12 +10,27 @@ except ImportError:
print("jaxdecomp not installed. Distributed functions will not work.")
distributed = False
from functools import partial
import jax
import jax.numpy as jnp
from jax._src import mesh as mesh_lib
from jax.experimental.shard_map import shard_map
from functools import partial
from jax.sharding import PartitionSpec as P
# NOTE
# This should not be used as a decorator
# Must be used inside a function only
# Example
# BAD
# @autoshmap
# def foo():
# pass
# GOOD
# def foo():
# return autoshmap(foo_impl)()
def autoshmap(f: Callable,
in_specs: Specs,
out_specs: Specs,
@ -34,31 +49,43 @@ def fft3d(x):
if distributed and not (mesh_lib.thread_resources.env.physical_mesh.empty):
return jaxdecomp.pfft3d(x.astype(jnp.complex64))
else:
return jnp.fft.rfftn(x)
return jnp.fft.fftn(x.astype(jnp.complex64))
def ifft3d(x):
if distributed and not (mesh_lib.thread_resources.env.physical_mesh.empty):
return jaxdecomp.pifft3d(x).real
else:
return jnp.fft.irfftn(x)
def halo_exchange(x):
if distributed and not (mesh_lib.thread_resources.env.physical_mesh.empty):
return jaxdecomp.halo_exchange(x)
return jnp.fft.ifftn(x).real
def get_halo_size(halo_size):
mesh = mesh_lib.thread_resources.env.physical_mesh
if mesh.empty:
zero_ext = (0, 0, 0)
zero_tuple = (0, 0)
return (zero_tuple, zero_tuple, zero_tuple), zero_ext
else:
pdims = mesh.devices.shape
halo_x = (0, 0) if pdims[0] == 1 else (halo_size, halo_size)
halo_y = (0, 0) if pdims[1] == 1 else (halo_size, halo_size)
halo_x_ext = 0 if pdims[0] == 1 else halo_size // 2
halo_y_ext = 0 if pdims[1] == 1 else halo_size // 2
return ((halo_x, halo_y, (0, 0)), (halo_x_ext, halo_y_ext, 0))
def halo_exchange(x, halo_extents, halo_periods=(True, True, True)):
mesh = mesh_lib.thread_resources.env.physical_mesh
if distributed and not (mesh.empty) and (halo_extents[0] > 0
or halo_extents[1] > 0):
return jaxdecomp.halo_exchange(x, halo_extents, halo_periods)
else:
return x
@partial(autoshmap,
in_specs=(P('x', 'y'), P()),
out_specs=P('x', 'y'))
def slice_pad_impl(x, pad_width):
return jnp.pad(x, pad_width)
@partial(autoshmap,
in_specs=(P('x', 'y'), P()),
out_specs=P('x', 'y'))
def slice_unpad_impl(x, pad_width):
halo_x, _ = pad_width[0]
halo_y, _ = pad_width[0]
@ -68,17 +95,28 @@ def slice_unpad_impl(x, pad_width):
# Apply corrections along y
x = x.at[:, halo_y:halo_y + halo_y // 2].add(x[:, :halo_y // 2])
x = x.at[:, -(halo_y + halo_y // 2):-halo_y].add(x[:, -halo_y // 2:])
return x
return x[halo_x:-halo_x, halo_y:-halo_y, :]
def slice_pad(x, pad_width):
if distributed and not (mesh_lib.thread_resources.env.physical_mesh.empty):
return slice_pad_impl(x, pad_width)
mesh = mesh_lib.thread_resources.env.physical_mesh
if distributed and not (mesh.empty) and (pad_width[0][0] > 0
or pad_width[1][0] > 0):
return autoshmap((partial(jnp.pad, pad_width=pad_width)),
in_specs=(P('x', 'y')),
out_specs=P('x', 'y'))(x)
else:
return x
def slice_unpad(x, pad_width):
if distributed and not (mesh_lib.thread_resources.env.physical_mesh.empty):
return slice_unpad_impl(x, pad_width)
mesh = mesh_lib.thread_resources.env.physical_mesh
if distributed and not (mesh.empty) and (pad_width[0][0] > 0
or pad_width[1][0] > 0):
return autoshmap(partial(slice_unpad_impl, pad_width=pad_width),
in_specs=(P('x', 'y')),
out_specs=P('x', 'y'))(x)
else:
return x