From c0efd849f2af1adf6b68dd34e8cd1e31bfa87eb6 Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Fri, 24 Sep 2021 13:50:28 -0500
Subject: [PATCH] return device scalar from nodal_sum/nodal_min/nodal_max
---
grudge/dt_utils.py | 16 ++++++++++------
grudge/models/__init__.py | 2 +-
grudge/reductions.py | 31 ++++++++++++-------------------
test/test_dt_utils.py | 3 ++-
test/test_grudge.py | 2 +-
test/test_op.py | 4 ++--
6 files changed, 28 insertions(+), 30 deletions(-)
diff --git a/grudge/dt_utils.py b/grudge/dt_utils.py
index 41031e39..57801476 100644
--- a/grudge/dt_utils.py
+++ b/grudge/dt_utils.py
@@ -179,12 +179,14 @@ def h_max_from_volume(
if dim is None:
dim = dcoll.dim
- ones = dcoll.discr_from_dd(dd).zeros(dcoll._setup_actx) + 1.0
- return nodal_max(
+ actx = dcoll._setup_actx
+
+ ones = dcoll.discr_from_dd(dd).zeros(actx) + 1.0
+ return actx.to_numpy(nodal_max(
dcoll,
dd,
elementwise_sum(dcoll, op.mass(dcoll, dd, ones))
- ) ** (1.0 / dim)
+ ) ** (1.0 / dim))
@memoize_on_first_arg
@@ -210,12 +212,14 @@ def h_min_from_volume(
if dim is None:
dim = dcoll.dim
- ones = dcoll.discr_from_dd(dd).zeros(dcoll._setup_actx) + 1.0
- return nodal_min(
+ actx = dcoll._setup_actx
+
+ ones = dcoll.discr_from_dd(dd).zeros(actx) + 1.0
+ return actx.to_numpy(nodal_min(
dcoll,
dd,
elementwise_sum(dcoll, op.mass(dcoll, dd, ones))
- ) ** (1.0 / dim)
+ ) ** (1.0 / dim))
def dt_geometric_factors(
diff --git a/grudge/models/__init__.py b/grudge/models/__init__.py
index 6c211c06..5e91f057 100644
--- a/grudge/models/__init__.py
+++ b/grudge/models/__init__.py
@@ -71,4 +71,4 @@ class HyperbolicOperator(Operator):
characteristic_lengthscales(actx, dcoll) / wavespeeds
)
- return op.nodal_min(dcoll, "vol", local_timesteps)
+ return actx.to_numpy(op.nodal_min(dcoll, "vol", local_timesteps))
diff --git a/grudge/reductions.py b/grudge/reductions.py
index 150ce01e..83192e42 100644
--- a/grudge/reductions.py
+++ b/grudge/reductions.py
@@ -144,7 +144,8 @@ def nodal_sum(dcoll: DiscretizationCollection, dd, vec) -> float:
from mpi4py import MPI
actx = vec.array_context
- return comm.allreduce(actx.to_numpy(nodal_sum_loc(dcoll, dd, vec)), op=MPI.SUM)
+ return actx.from_numpy(
+ comm.allreduce(actx.to_numpy(nodal_sum_loc(dcoll, dd, vec)), op=MPI.SUM))
def nodal_sum_loc(dcoll: DiscretizationCollection, dd, vec) -> float:
@@ -159,7 +160,6 @@ def nodal_sum_loc(dcoll: DiscretizationCollection, dd, vec) -> float:
return sum(nodal_sum_loc(dcoll, dd, vec[idx])
for idx in np.ndindex(vec.shape))
- # FIXME: do not force result onto host
actx = vec.array_context
return sum([actx.np.sum(grp_ary) for grp_ary in vec])
@@ -178,8 +178,10 @@ def nodal_min(dcoll: DiscretizationCollection, dd, vec) -> float:
# NOTE: Don't move this
from mpi4py import MPI
+ actx = vec.array_context
- return comm.allreduce(nodal_min_loc(dcoll, dd, vec), op=MPI.MIN)
+ return actx.from_numpy(
+ comm.allreduce(actx.to_numpy(nodal_min_loc(dcoll, dd, vec)), op=MPI.MIN))
def nodal_min_loc(dcoll: DiscretizationCollection, dd, vec) -> float:
@@ -196,16 +198,10 @@ def nodal_min_loc(dcoll: DiscretizationCollection, dd, vec) -> float:
for idx in np.ndindex(vec.shape))
actx = vec.array_context
-
- # FIXME: do not force result onto host
- # Host transfer is needed for now because actx.np.minimum does not succeed
- # on array scalars.
- # https://github.com/inducer/arraycontext/issues/49#issuecomment-869266944
return reduce(
lambda acc, grp_ary: actx.np.minimum(
- acc,
- actx.to_numpy(actx.np.min(grp_ary))[()]),
- vec, np.inf)
+ acc, actx.np.min(grp_ary)),
+ vec, actx.from_numpy(np.array(np.inf)))
def nodal_max(dcoll: DiscretizationCollection, dd, vec) -> float:
@@ -222,8 +218,10 @@ def nodal_max(dcoll: DiscretizationCollection, dd, vec) -> float:
# NOTE: Don't move this
from mpi4py import MPI
+ actx = vec.array_context
- return comm.allreduce(nodal_max_loc(dcoll, dd, vec), op=MPI.MAX)
+ return actx.from_numpy(
+ comm.allreduce(actx.to_numpy(nodal_max_loc(dcoll, dd, vec)), op=MPI.MAX))
def nodal_max_loc(dcoll: DiscretizationCollection, dd, vec) -> float:
@@ -241,15 +239,10 @@ def nodal_max_loc(dcoll: DiscretizationCollection, dd, vec) -> float:
actx = vec.array_context
- # FIXME: do not force result onto host
- # Host transfer is needed for now because actx.np.minimum does not succeed
- # on array scalars.
- # https://github.com/inducer/arraycontext/issues/49#issuecomment-869266944
return reduce(
lambda acc, grp_ary: actx.np.maximum(
- acc,
- actx.to_numpy(actx.np.max(grp_ary))[()]),
- vec, -np.inf)
+ acc, actx.np.max(grp_ary)),
+ vec, actx.from_numpy(np.array(-np.inf)))
def integral(dcoll: DiscretizationCollection, dd, vec) -> float:
diff --git a/test/test_dt_utils.py b/test/test_dt_utils.py
index f3f3f9a5..9d4c345d 100644
--- a/test/test_dt_utils.py
+++ b/test/test_dt_utils.py
@@ -70,7 +70,8 @@ def test_geometric_factors_regular_refinement(actx_factory, name):
mesh = builder.get_mesh(resolution, builder.mesh_order)
dcoll = DiscretizationCollection(actx, mesh, order=builder.order)
min_factors.append(
- op.nodal_min(dcoll, "vol", thaw(dt_geometric_factors(dcoll), actx))
+ actx.to_numpy(
+ op.nodal_min(dcoll, "vol", thaw(dt_geometric_factors(dcoll), actx)))
)
# Resolution is doubled each refinement, so the ratio of consecutive
diff --git a/test/test_grudge.py b/test/test_grudge.py
index f6e4ad9b..43c3535d 100644
--- a/test/test_grudge.py
+++ b/test/test_grudge.py
@@ -637,7 +637,7 @@ def test_surface_divergence_theorem(actx_factory, mesh_name, visualize=False):
h_max = h_max_from_volume(dcoll)
eoc_global.add_data_point(h_max, actx.to_numpy(err_global))
- eoc_local.add_data_point(h_max, err_local)
+ eoc_local.add_data_point(h_max, actx.to_numpy(err_local))
if visualize:
from grudge.shortcuts import make_visualizer
diff --git a/test/test_op.py b/test/test_op.py
index 92c08889..364a988f 100644
--- a/test/test_op.py
+++ b/test/test_op.py
@@ -151,7 +151,7 @@ def test_gradient(actx_factory, form, dim, order, vectorize, nested,
("expected_grad_u", expected_grad_u),
], overwrite=True)
- rel_linf_err = (
+ rel_linf_err = actx.to_numpy(
op.norm(dcoll, grad_u - expected_grad_u, np.inf)
/ op.norm(dcoll, expected_grad_u, np.inf))
eoc_rec.add_data_point(1./n, rel_linf_err)
@@ -267,7 +267,7 @@ def test_divergence(actx_factory, form, dim, order, vectorize, nested,
("expected_div_u", expected_div_u),
], overwrite=True)
- rel_linf_err = (
+ rel_linf_err = actx.to_numpy(
op.norm(dcoll, div_u - expected_div_u, np.inf)
/ op.norm(dcoll, expected_div_u, np.inf))
eoc_rec.add_data_point(1./n, rel_linf_err)
--
GitLab