diff --git a/examples/ablation-workshop.py b/examples/ablation-workshop.py index 8fc6c12f6..1a7dd77b0 100644 --- a/examples/ablation-workshop.py +++ b/examples/ablation-workshop.py @@ -1212,13 +1212,16 @@ def my_post_step(step, t, dt, state): dest="restart_file", nargs="?", action="store", help="simulation restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn warn("Automatically turning off DV logging. MIRGE-Com Issue(578)") from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/advection_diffusion_reaction.py b/examples/advection_diffusion_reaction.py index 145302383..2b74b87b2 100644 --- a/examples/advection_diffusion_reaction.py +++ b/examples/advection_diffusion_reaction.py @@ -331,13 +331,16 @@ def my_rhs(t, u): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() lazy = args.lazy from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class(lazy=args.lazy, distributed=True, profiling=args.profiling, - numpy=args.numpy) + numpy=args.numpy, + cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/autoignition.py b/examples/autoignition.py index 2ec92e9e1..9a47fc702 100644 --- a/examples/autoignition.py +++ b/examples/autoignition.py @@ -657,6 +657,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn warn("Automatically turning off DV logging. MIRGE-Com Issue(578)") @@ -672,7 +674,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/blasius.py b/examples/blasius.py index 1eb13ad2f..40f417c81 100644 --- a/examples/blasius.py +++ b/examples/blasius.py @@ -523,6 +523,8 @@ def my_post_step(step, t, dt, state): help="enable lazy evaluation [OFF]") parser.add_argument("--numpy", action="store_true", help="use numpy-based eager actx.") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() @@ -538,7 +540,8 @@ def my_post_step(step, t, dt, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/combozzle.py b/examples/combozzle.py index 14eabf8cf..77e394efc 100644 --- a/examples/combozzle.py +++ b/examples/combozzle.py @@ -1255,6 +1255,8 @@ def dummy_rhs(t, state): parser.add_argument("--casename", help="casename to use for i/o") parser.add_argument("--tpe", action="store_true", help="Use tensor product elements (quads/hexes).") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -1271,7 +1273,8 @@ def dummy_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/doublemach.py b/examples/doublemach.py index b5c9df3fe..42ff2ac56 100644 --- a/examples/doublemach.py +++ b/examples/doublemach.py @@ -443,6 +443,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -455,7 +457,7 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class(lazy=args.lazy, distributed=True, profiling=args.profiling, - numpy=args.numpy) + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/doublemach_physical_av.py b/examples/doublemach_physical_av.py index fac0879fd..ef10d5334 100644 --- a/examples/doublemach_physical_av.py +++ b/examples/doublemach_physical_av.py @@ -716,6 +716,8 @@ def _my_rhs_phys_visc_div_av(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -730,7 +732,7 @@ def _my_rhs_phys_visc_div_av(t, state): actx_class = get_reasonable_array_context_class(lazy=args.lazy, distributed=True, profiling=args.profiling, - numpy=args.numpy) + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/heat-source.py b/examples/heat-source.py index db5970dd9..9d3c066df 100644 --- a/examples/heat-source.py +++ b/examples/heat-source.py @@ -245,11 +245,14 @@ def my_post_step(step, t, dt, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/hotplate.py b/examples/hotplate.py index ca3444505..71bc21dd7 100644 --- a/examples/hotplate.py +++ b/examples/hotplate.py @@ -448,6 +448,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -460,7 +462,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/lump.py b/examples/lump.py index 0891d138b..eb19ffd20 100644 --- a/examples/lump.py +++ b/examples/lump.py @@ -383,6 +383,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -395,7 +397,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/mixture.py b/examples/mixture.py index f6276615d..b4c8a11bb 100644 --- a/examples/mixture.py +++ b/examples/mixture.py @@ -447,6 +447,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn warn("Automatically turning off DV logging. MIRGE-Com Issue(578)") @@ -461,7 +463,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/multiple-volumes.py b/examples/multiple-volumes.py index 2cc13f873..9eb1fc369 100644 --- a/examples/multiple-volumes.py +++ b/examples/multiple-volumes.py @@ -392,6 +392,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -404,7 +406,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/orthotropic-diffusion.py b/examples/orthotropic-diffusion.py index 0014ad059..f48880a13 100644 --- a/examples/orthotropic-diffusion.py +++ b/examples/orthotropic-diffusion.py @@ -204,11 +204,14 @@ def my_post_step(step, t, dt, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/poiseuille-multispecies.py b/examples/poiseuille-multispecies.py index 8d32aeea2..1caf71056 100644 --- a/examples/poiseuille-multispecies.py +++ b/examples/poiseuille-multispecies.py @@ -486,6 +486,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -498,7 +500,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/poiseuille.py b/examples/poiseuille.py index c8f2f2a98..7200796c5 100644 --- a/examples/poiseuille.py +++ b/examples/poiseuille.py @@ -475,6 +475,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -487,7 +489,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/pulse-mixture.py b/examples/pulse-mixture.py index 56d068141..dc0e97426 100644 --- a/examples/pulse-mixture.py +++ b/examples/pulse-mixture.py @@ -421,6 +421,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -433,7 +435,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/pulse-tpe.py b/examples/pulse-tpe.py index 50d10e84f..218d7d0e9 100644 --- a/examples/pulse-tpe.py +++ b/examples/pulse-tpe.py @@ -348,6 +348,9 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") + args = parser.parse_args() from warnings import warn @@ -362,7 +365,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/pulse.py b/examples/pulse.py index 74b5bb6aa..f6e84206c 100644 --- a/examples/pulse.py +++ b/examples/pulse.py @@ -342,6 +342,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -354,7 +356,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/scalar-advdiff.py b/examples/scalar-advdiff.py index 966252c4c..e5a7bb5bf 100644 --- a/examples/scalar-advdiff.py +++ b/examples/scalar-advdiff.py @@ -467,6 +467,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -480,7 +482,7 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class(lazy=args.lazy, distributed=True, profiling=args.profiling, - numpy=args.numpy) + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/scalar-lump.py b/examples/scalar-lump.py index 8716bc2c7..46afa6eac 100644 --- a/examples/scalar-lump.py +++ b/examples/scalar-lump.py @@ -393,6 +393,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -405,7 +407,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/sod.py b/examples/sod.py index ab315a0c5..91713da11 100644 --- a/examples/sod.py +++ b/examples/sod.py @@ -391,6 +391,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -403,7 +405,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/taylor-green.py b/examples/taylor-green.py index 4ce8a04df..062004b82 100644 --- a/examples/taylor-green.py +++ b/examples/taylor-green.py @@ -333,6 +333,8 @@ def my_rhs(t, state): help="turn on logging") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -345,7 +347,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/thermally-coupled.py b/examples/thermally-coupled.py index 93d737bde..e3e91561f 100644 --- a/examples/thermally-coupled.py +++ b/examples/thermally-coupled.py @@ -591,6 +591,8 @@ def my_rhs_and_gradients(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -603,7 +605,8 @@ def my_rhs_and_gradients(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/vortex.py b/examples/vortex.py index 26782c872..3763227fd 100644 --- a/examples/vortex.py +++ b/examples/vortex.py @@ -401,6 +401,8 @@ def my_rhs(t, state): help="use numpy-based eager actx.") parser.add_argument("--restart_file", help="root name of restart file") parser.add_argument("--casename", help="casename to use for i/o") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() from warnings import warn @@ -413,7 +415,8 @@ def my_rhs(t, state): from mirgecom.array_context import get_reasonable_array_context_class actx_class = get_reasonable_array_context_class( - lazy=args.lazy, distributed=True, profiling=args.profiling, numpy=args.numpy) + lazy=args.lazy, distributed=True, profiling=args.profiling, + numpy=args.numpy, cupy=args.cupy) logging.basicConfig(format="%(message)s", level=logging.INFO) if args.casename: diff --git a/examples/wave.py b/examples/wave.py index c0c1e2d70..dc6f6e51a 100644 --- a/examples/wave.py +++ b/examples/wave.py @@ -271,6 +271,8 @@ def rhs(t, w): help="use numpy-based eager actx.") parser.add_argument("--mpi", default=True, action=argparse.BooleanOptionalAction, help="use MPI") + parser.add_argument("--cupy", action="store_true", + help="use cupy-based eager actx.") args = parser.parse_args() casename = args.casename or "wave" @@ -278,7 +280,8 @@ def rhs(t, w): actx_class = get_reasonable_array_context_class(lazy=args.lazy, distributed=args.mpi, profiling=args.profiling, - numpy=args.numpy) + numpy=args.numpy, + cupy=args.cupy) if args.mpi: main_func = main diff --git a/mirgecom/array_context.py b/mirgecom/array_context.py index b94e93c77..29d4ed39e 100644 --- a/mirgecom/array_context.py +++ b/mirgecom/array_context.py @@ -45,11 +45,33 @@ def get_reasonable_array_context_class(*, lazy: bool, distributed: bool, - profiling: bool, numpy: bool = False) -> Type[ArrayContext]: + profiling: bool, numpy: bool = False, + cupy: bool = False) -> Type[ArrayContext]: """Return a :class:`~arraycontext.ArrayContext` with the given constraints.""" if lazy and profiling: raise ValueError("Can't specify both lazy and profiling") + if numpy and cupy: + raise ValueError("Can't specify both numpy and cupy") + + if cupy: + if profiling: + raise ValueError("Can't specify both cupy and profiling") + if lazy: + raise ValueError("Can't specify both cupy and lazy") + + from warnings import warn + warn("The CupyArrayContext is still under development") + + if distributed: + from grudge.array_context import ( # type: ignore[attr-defined] # pylint: disable=no-name-in-module # noqa: E501 + MPICupyArrayContext) + return MPICupyArrayContext + else: + from grudge.array_context import ( # type: ignore[attr-defined] # pylint: disable=no-name-in-module # noqa: E501 + CupyArrayContext) + return CupyArrayContext + if numpy: if profiling: raise ValueError("Can't specify both numpy and profiling") @@ -100,6 +122,18 @@ def actx_class_is_numpy(actx_class: Type[ArrayContext]) -> bool: return issubclass(actx_class, NumpyArrayContext) +def actx_class_is_cupy(actx_class: Type[ArrayContext]) -> bool: + """Return True if *actx_class* is cupy-based.""" + try: + from grudge.array_context import CupyArrayContext # type: ignore[attr-defined] # noqa: E501 + if issubclass(actx_class, CupyArrayContext): + return True + else: + return False + except ImportError: + return False + + def actx_class_is_distributed(actx_class: Type[ArrayContext]) -> bool: """Return True if *actx_class* is distributed.""" from grudge.array_context import MPIBasedArrayContext @@ -301,13 +335,23 @@ def initialize_actx( if comm: actx_kwargs["mpi_communicator"] = comm - if actx_class_is_numpy(actx_class): - from grudge.array_context import MPINumpyArrayContext - if comm: - assert issubclass(actx_class, MPINumpyArrayContext) + # Special handling for NumpyArrayContext/CupyArrayContext + # since they need no CL context + if actx_class_is_numpy(actx_class) or actx_class_is_cupy(actx_class): + if actx_class_is_numpy(actx_class): + from grudge.array_context import MPINumpyArrayContext + if comm: + assert issubclass(actx_class, MPINumpyArrayContext) + else: + assert not issubclass(actx_class, MPINumpyArrayContext) else: - assert not issubclass(actx_class, MPINumpyArrayContext) + from grudge.array_context import MPICupyArrayContext # type: ignore[attr-defined] # pylint: disable=no-name-in-module # noqa: E501 + if comm: + assert issubclass(actx_class, MPICupyArrayContext) + else: + assert not issubclass(actx_class, MPICupyArrayContext) else: + # PyOpenCL-based actx cl_ctx = cl.create_some_context() if actx_class_is_profiling(actx_class): queue = cl.CommandQueue(cl_ctx, diff --git a/mirgecom/logging_quantities.py b/mirgecom/logging_quantities.py index 67cf65527..33631d76d 100644 --- a/mirgecom/logging_quantities.py +++ b/mirgecom/logging_quantities.py @@ -78,7 +78,7 @@ def initialize_logmgr(enable_logmgr: bool, logmgr.enable_save_on_sigterm() add_run_info(logmgr) - add_package_versions(logmgr) + # add_package_versions(logmgr) add_general_quantities(logmgr) add_simulation_quantities(logmgr)