From 0d34e275fb34b8257d5847ee597596a726d2c964 Mon Sep 17 00:00:00 2001 From: Tuomas Karna Date: Thu, 7 Nov 2024 21:28:20 +0200 Subject: [PATCH 1/4] wave-equation: remove fromfunction initialization --- examples/wave_equation.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/examples/wave_equation.py b/examples/wave_equation.py index b9dc507..e95fd0c 100644 --- a/examples/wave_equation.py +++ b/examples/wave_equation.py @@ -54,18 +54,18 @@ def run(n, backend, datatype, benchmark_mode): if backend == "sharpy": import sharpy as np from sharpy import fini, init, sync - from sharpy.numpy import fromfunction as _fromfunction device = os.getenv("SHARPY_DEVICE", "") create_full = partial(np.full, device=device) - fromfunction = partial(_fromfunction, device=device) + + def transpose(a): + return np.permute_dims(a, [1, 0]) all_axes = [0, 1] init(False) elif backend == "numpy": import numpy as np - from numpy import fromfunction if comm is not None: assert ( @@ -73,6 +73,7 @@ def run(n, backend, datatype, benchmark_mode): ), "Numpy backend only supports serial execution." create_full = np.full + transpose = np.transpose fini = sync = lambda x=None: None all_axes = None @@ -110,17 +111,23 @@ def run(n, backend, datatype, benchmark_mode): t_export = 0.02 t_end = 1.0 - # coordinate arrays - x_t_2d = fromfunction( - lambda i, j: xmin + i * dx + dx / 2, (nx, ny), dtype=dtype - ) - y_t_2d = fromfunction( - lambda i, j: ymin + j * dy + dy / 2, (nx, ny), dtype=dtype - ) + def ind_arr(shape, columns=False): + """Construct an (nx, ny) array where each row/col is an arange""" + nx, ny = shape + if columns: + ind = np.arange(0, nx * ny, 1, dtype=np.int32) % nx + ind = transpose(np.reshape(ind, (ny, nx))) + else: + ind = np.arange(0, nx * ny, 1, dtype=np.int32) % ny + ind = np.reshape(ind, (nx, ny)) + return ind.astype(dtype) + # coordinate arrays T_shape = (nx, ny) U_shape = (nx + 1, ny) V_shape = (nx, ny + 1) + x_t_2d = xmin + ind_arr(T_shape, True) * dx + dx / 2 + y_t_2d = ymin + ind_arr(T_shape) * dy + dy / 2 dofs_T = int(numpy.prod(numpy.asarray(T_shape))) dofs_U = int(numpy.prod(numpy.asarray(U_shape))) From 7fb1742dcb18577fce10959b7540e1def70f1b80 Mon Sep 17 00:00:00 2001 From: Tuomas Karna Date: Thu, 7 Nov 2024 21:28:43 +0200 Subject: [PATCH 2/4] shallow-water: remove fromfunction initialization --- examples/shallow_water.py | 51 +++++++++++++++++++-------------------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/examples/shallow_water.py b/examples/shallow_water.py index 06ea67b..3c8118d 100644 --- a/examples/shallow_water.py +++ b/examples/shallow_water.py @@ -54,18 +54,18 @@ def run(n, backend, datatype, benchmark_mode): if backend == "sharpy": import sharpy as np from sharpy import fini, init, sync - from sharpy.numpy import fromfunction as _fromfunction device = os.getenv("SHARPY_DEVICE", "") create_full = partial(np.full, device=device) - fromfunction = partial(_fromfunction, device=device) + + def transpose(a): + return np.permute_dims(a, [1, 0]) all_axes = [0, 1] init(False) elif backend == "numpy": import numpy as np - from numpy import fromfunction if comm is not None: assert ( @@ -73,6 +73,7 @@ def run(n, backend, datatype, benchmark_mode): ), "Numpy backend only supports serial execution." create_full = np.full + transpose = np.transpose fini = sync = lambda x=None: None all_axes = None @@ -110,34 +111,32 @@ def run(n, backend, datatype, benchmark_mode): t_export = 0.02 t_end = 1.0 - # coordinate arrays - x_t_2d = fromfunction( - lambda i, j: xmin + i * dx + dx / 2, - (nx, ny), - dtype=dtype, - ) - y_t_2d = fromfunction( - lambda i, j: ymin + j * dy + dy / 2, - (nx, ny), - dtype=dtype, - ) - x_u_2d = fromfunction(lambda i, j: xmin + i * dx, (nx + 1, ny), dtype=dtype) - y_u_2d = fromfunction( - lambda i, j: ymin + j * dy + dy / 2, - (nx + 1, ny), - dtype=dtype, - ) - x_v_2d = fromfunction( - lambda i, j: xmin + i * dx + dx / 2, - (nx, ny + 1), - dtype=dtype, - ) - y_v_2d = fromfunction(lambda i, j: ymin + j * dy, (nx, ny + 1), dtype=dtype) + def ind_arr(shape, columns=False): + """Construct an (nx, ny) array where each row/col is an arange""" + nx, ny = shape + if columns: + ind = np.arange(0, nx * ny, 1, dtype=np.int32) % nx + ind = transpose(np.reshape(ind, (ny, nx))) + else: + ind = np.arange(0, nx * ny, 1, dtype=np.int32) % ny + ind = np.reshape(ind, (nx, ny)) + return ind.astype(dtype) + # coordinate arrays T_shape = (nx, ny) U_shape = (nx + 1, ny) V_shape = (nx, ny + 1) F_shape = (nx + 1, ny + 1) + sync() + x_t_2d = xmin + ind_arr(T_shape, True) * dx + dx / 2 + y_t_2d = ymin + ind_arr(T_shape) * dy + dy / 2 + + x_u_2d = xmin + ind_arr(U_shape, True) * dx + y_u_2d = ymin + ind_arr(U_shape) * dy + dy / 2 + + x_v_2d = xmin + ind_arr(V_shape, True) * dx + dx / 2 + y_v_2d = ymin + ind_arr(V_shape) * dy + sync() dofs_T = int(numpy.prod(numpy.asarray(T_shape))) dofs_U = int(numpy.prod(numpy.asarray(U_shape))) From 048f1c68115494c5b7fad8a9378c978f263b90ef Mon Sep 17 00:00:00 2001 From: Tuomas Karna Date: Thu, 7 Nov 2024 21:46:09 +0200 Subject: [PATCH 3/4] wave-equation: exclude jit compilation from timing --- examples/wave_equation.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/examples/wave_equation.py b/examples/wave_equation.py index e95fd0c..defd2c6 100644 --- a/examples/wave_equation.py +++ b/examples/wave_equation.py @@ -169,9 +169,6 @@ def exact_elev(t, x_t_2d, y_t_2d, lx, ly): sol_t = numpy.cos(2 * omega * t) return amp * sol_x * sol_y * sol_t - # inital elevation - e[:, :] = exact_elev(0.0, x_t_2d, y_t_2d, lx, ly) - # compute time step alpha = 0.5 c = (g * h) ** 0.5 @@ -222,6 +219,16 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2): v[:, 1:-1] = v[:, 1:-1] / 3.0 + 2.0 / 3.0 * (v2[:, 1:-1] + dt * dvdt) e[:, :] = e[:, :] / 3.0 + 2.0 / 3.0 * (e2[:, :] + dt * dedt) + # warm up jit cache + step(u, v, e, u1, v1, e1, u2, v2, e2) + sync() + + # initial solution + e[:, :] = exact_elev(0.0, x_t_2d, y_t_2d, lx, ly) + u[:, :] = create_full(U_shape, 0.0, dtype) + v[:, :] = create_full(V_shape, 0.0, dtype) + sync() + t = 0 i_export = 0 next_t_export = 0 From aba4ce814ef0d021884f6dafb5568b1dea2e553f Mon Sep 17 00:00:00 2001 From: Tuomas Karna Date: Thu, 7 Nov 2024 22:08:05 +0200 Subject: [PATCH 4/4] shallow-water: exclude jit compilation from timing --- examples/shallow_water.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/examples/shallow_water.py b/examples/shallow_water.py index 3c8118d..a5798a4 100644 --- a/examples/shallow_water.py +++ b/examples/shallow_water.py @@ -204,14 +204,6 @@ def bathymetry(x_t_2d, y_t_2d, lx, ly): bath = 1.0 return bath * create_full(T_shape, 1.0, dtype) - # inital elevation - u0, v0, e0 = exact_solution( - 0, x_t_2d, y_t_2d, x_u_2d, y_u_2d, x_v_2d, y_v_2d - ) - e[:, :] = e0 - u[:, :] = u0 - v[:, :] = v0 - # set bathymetry h[:, :] = bathymetry(x_t_2d, y_t_2d, lx, ly) # steady state potential energy @@ -328,6 +320,18 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2): v[:, 1:-1] = v[:, 1:-1] / 3.0 + 2.0 / 3.0 * (v2[:, 1:-1] + dt * dvdt) e[:, :] = e[:, :] / 3.0 + 2.0 / 3.0 * (e2[:, :] + dt * dedt) + # warm up jit cache + step(u, v, e, u1, v1, e1, u2, v2, e2) + sync() + + # initial solution + u0, v0, e0 = exact_solution( + 0, x_t_2d, y_t_2d, x_u_2d, y_u_2d, x_v_2d, y_v_2d + ) + e[:, :] = e0 + u[:, :] = u0 + v[:, :] = v0 + t = 0 i_export = 0 next_t_export = 0