Skip to content

Commit

Permalink
Merge branch 'topic/default/jit2boost' into 'branch/default'
Browse files Browse the repository at this point in the history
Replace transonic.jit by boost

See merge request fluiddyn/fluidsim!348
  • Loading branch information
paugier committed Jan 2, 2024
2 parents 7de523b + c5f391d commit 7322f7a
Show file tree
Hide file tree
Showing 20 changed files with 337 additions and 301 deletions.
2 changes: 2 additions & 0 deletions .hgignore
Original file line number Diff line number Diff line change
Expand Up @@ -83,3 +83,5 @@ scratch
fluidsim/_version.py

doc/examples/htmlcov/*

fluidsim/build_conf.txt
6 changes: 3 additions & 3 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ MPI_NUM_PROCS ?= 2
develop: develop_lib
pip install meson-python ninja numpy "pythran>=0.9.7"
pip install --force-reinstall transonic@hg+https://foss.heptapod.net/fluiddyn/transonic
pip install -e .[dev] --no-build-isolation
pip install -e ".[dev]" --no-build-isolation

develop_lib:
cd lib && pip install -e .
Expand Down Expand Up @@ -41,10 +41,10 @@ shortlog:
@hg log -M -r$(RELEASE): --template '- {desc|firstline} (:rev:`{node|short}`)\n'

black:
black -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|doc/_build|\.ipynb_checkpoints/*)/"
black -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"

black_check:
black --check -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|doc/_build|\.ipynb_checkpoints/*)/"
black --check -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"

tests:
pytest -v lib
Expand Down
6 changes: 0 additions & 6 deletions doc/install.rst
Original file line number Diff line number Diff line change
Expand Up @@ -100,12 +100,6 @@ Fluidsim is also sensitive to the environment variables:

- ``FLUIDDYN_PATH_SCRATCH``: working directory (can be useful on some clusters).

- ``TRANSONIC_COMPILE_JIT``: set this variable to force JIT compilation using
``transonic`` while running tests. This is not necessary, but could be useful
for troubleshooting if simulations freeze. For example::

TRANSONIC_COMPILE_JIT=1 fluidsim-test -m fluidsim.solvers.sw1l

Warning about re-installing fluidsim and fluidfft with new build options
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Expand Down
13 changes: 0 additions & 13 deletions fluidsim/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,19 +55,6 @@ def _show(*args, **kwargs):

plt.show = _show

if all(
env_var not in os.environ
for env_var in ("FLUID_COMPILE_CACHEDJIT", "TRANSONIC_COMPILE_JIT")
):
mpi.printby0("Compilation of jit functions disabled.")
from transonic import set_compile_jit

set_compile_jit(False)
elif "FLUID_COMPILE_CACHEDJIT" in os.environ:
mpi.printby0(
"WARNING: FLUID_COMPILE_CACHEDJIT is deprecated, use "
"TRANSONIC_COMPILE_JIT instead."
)

from ._version import __version__, get_local_version

Expand Down
6 changes: 6 additions & 0 deletions fluidsim/solvers/ns2d/bouss/meson.build
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,9 @@ py.install_sources(
python_sources,
subdir: 'fluidsim/solvers/ns2d/bouss',
)

run_command(['transonic', '--meson', '--backend', backend, 'solver.py'], check: true)

foreach be : backends
subdir('__' + be + '__')
endforeach
4 changes: 2 additions & 2 deletions fluidsim/solvers/ns2d/bouss/solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
"""
import numpy as np

from transonic import jit, Array
from transonic import boost, Array

from fluidsim.base.setofvariables import SetOfVariables
from fluidsim.solvers.ns2d.solver import InfoSolverNS2D, Simul as SimulNS2D
Expand All @@ -17,7 +17,7 @@
AF = Array[np.float64, "2d"]


@jit
@boost
def tendencies_nonlin_ns2dbouss(
ux: AF, uy: AF, px_rot: AF, py_rot: AF, px_b: AF, py_b: AF
):
Expand Down
6 changes: 6 additions & 0 deletions fluidsim/solvers/ns2d/strat/meson.build
Original file line number Diff line number Diff line change
Expand Up @@ -14,3 +14,9 @@ py.install_sources(
)

subdir('output')

run_command(['transonic', '--meson', '--backend', backend, 'solver.py'], check: true)

foreach be : backends
subdir('__' + be + '__')
endforeach
4 changes: 2 additions & 2 deletions fluidsim/solvers/ns2d/strat/solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@

import numpy as np

from transonic import jit, Array
from transonic import boost, Array

from fluidsim.base.setofvariables import SetOfVariables

Expand All @@ -18,7 +18,7 @@
AF = Array[np.float64, "2d"]


@jit
@boost
def tendencies_nonlin_ns2dstrat(
ux: AF, uy: AF, px_rot: AF, py_rot: AF, px_b: AF, py_b: AF, N: float
):
Expand Down
6 changes: 6 additions & 0 deletions fluidsim/solvers/plate2d/meson.build
Original file line number Diff line number Diff line change
Expand Up @@ -16,3 +16,9 @@ py.install_sources(
)

subdir('output')

run_command(['transonic', '--meson', '--backend', backend, 'operators.py'], check: true)

foreach be : backends
subdir('__' + be + '__')
endforeach
6 changes: 3 additions & 3 deletions fluidsim/solvers/plate2d/operators.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,15 @@

import numpy as np

from transonic import jit, Array
from transonic import boost, Array

from ...operators.operators2d import OperatorsPseudoSpectral2D

AC2 = Array[np.complex128, "2d"]
AF2 = Array[np.float64, "2d"]


@jit
@boost
def monge_ampere_step0(a_fft: AC2, b_fft: AC2, KX2: AF2, KY2: AF2, KXKZ: AF2):
pxx_a_fft = -a_fft * KX2
pyy_a_fft = -a_fft * KY2
Expand All @@ -30,7 +30,7 @@ def monge_ampere_step0(a_fft: AC2, b_fft: AC2, KX2: AF2, KY2: AF2, KXKZ: AF2):
return pxx_a_fft, pyy_a_fft, pxy_a_fft, pxx_b_fft, pyy_b_fft, pxy_b_fft


@jit
@boost
def monge_ampere_step1(
pxx_a: AF2, pyy_a: AF2, pxy_a: AF2, pxx_b: AF2, pyy_b: AF2, pxy_b: AF2
):
Expand Down
6 changes: 3 additions & 3 deletions fluidsim/solvers/plate2d/output/correlations_freq.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,14 @@
import h5py
import matplotlib.pyplot as plt

from transonic import jit
from transonic import boost
from fluiddyn.util import mpi
from fluiddyn.calcul.easypyfft import FFTW1DReal2Complex

from fluidsim.base.output.base import SpecificOutput


@jit
@boost
def compute_correl4_seq(
q_fftt: "complex128[][]", iomegas1: "int32[]", nb_omegas: int, nb_xs_seq: int
):
Expand Down Expand Up @@ -94,7 +94,7 @@ def compute_correl4_seq(
return corr4


@jit
@boost
def compute_correl2_seq(
q_fftt: "complex128[][]", iomegas1: "int32[]", nb_omegas: int, nb_xs_seq: int
):
Expand Down
6 changes: 6 additions & 0 deletions fluidsim/solvers/plate2d/output/meson.build
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,9 @@ py.install_sources(
python_sources,
subdir: 'fluidsim/solvers/plate2d/output'
)

run_command(['transonic', '--meson', '--backend', backend, 'correlations_freq.py'], check: true)

foreach be : backends
subdir('__' + be + '__')
endforeach
5 changes: 5 additions & 0 deletions fluidsim/solvers/sphere/sw1l/meson.build
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,8 @@ py.install_sources(
subdir: 'fluidsim/solvers/sphere/sw1l'
)

run_command(['transonic', '--meson', '--backend', backend, 'solver.py'], check: true)

foreach be : backends
subdir('__' + be + '__')
endforeach
4 changes: 2 additions & 2 deletions fluidsim/solvers/sphere/sw1l/solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
:private-members:
"""
from transonic import jit
from transonic import boost
from fluidsim.base.sphericalharmo.solver import (
InfoSolverSphericalHarmo,
SimulSphericalHarmo,
Expand All @@ -22,7 +22,7 @@
Af = "float64[:, :]"


@jit
@boost
def compute_Frot(rot: Af, ux: Af, uy: Af, f_radial: Af):
"""Compute cross-product of absolute potential vorticity with velocity."""
rot_abs = rot + f_radial
Expand Down
9 changes: 9 additions & 0 deletions fluidsim/solvers/sw1l/meson.build
Original file line number Diff line number Diff line change
Expand Up @@ -18,3 +18,12 @@ subdir('exactlin')
subdir('modified')
subdir('onlywaves')
subdir('output')

run_command(
['transonic', '--meson', '--backend', backend, 'operators.py', 'solver.py'],
check: true
)

foreach be : backends
subdir('__' + be + '__')
endforeach
6 changes: 3 additions & 3 deletions fluidsim/solvers/sw1l/operators.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import numpy as np
from fluiddyn.util.compat import cached_property

from transonic import boost, jit, Array
from transonic import boost, Array

from fluidsim.operators.operators2d import (
OperatorsPseudoSpectral2D,
Expand All @@ -25,7 +25,7 @@
AF = Array[np.float64, "2d"]


@jit
@boost
def _qapamfft_from_uxuyetafft(
ux_fft: AC,
uy_fft: AC,
Expand Down Expand Up @@ -309,7 +309,7 @@ def apamfft_from_adfft(self, a_fft, d_fft):
am_fft = 0.5 * (a_fft - Delta_a_fft)
return ap_fft, am_fft

@jit
@boost
def divfft_from_apamfft(self, ap_fft: AC, am_fft: AC):
"""Return div from the eigen modes ap and am."""
n0 = self.nK0_loc
Expand Down
Loading

0 comments on commit 7322f7a

Please sign in to comment.