Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use cupy when available with the particle container wrapper #4203

Merged
merged 10 commits into from
Sep 23, 2023
35 changes: 35 additions & 0 deletions Python/pywarpx/LoadThirdParty.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
# This file is part of WarpX.
#
# Authors: Axel Huebl
# License: BSD-3-Clause-LBNL

from ._libwarpx import libwarpx


def load_cupy():
"""
This is a helper for
https://docs.cupy.dev/en/stable/user_guide/basic.html
"""
amr = libwarpx.amr
status = None

if amr.Config.have_gpu:
try:
import cupy as cp
xp = cp
# Note: found and will use cupy
except ImportError:
status = "Warning: GPU found but cupy not available! Trying managed memory in numpy..."
import numpy as np
xp = np
if amr.Config.gpu_backend == "SYCL":
status = "Warning: SYCL GPU backend not yet implemented for Python"
import numpy as np
xp = np

else:
import numpy as np
xp = np
# Note: found and will use numpy
return xp, status
1 change: 1 addition & 0 deletions Python/pywarpx/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
from .HybridPICModel import hybridpicmodel
from .Interpolation import interpolation
from .Lasers import lasers
from .LoadThirdParty import load_cupy
from .PSATD import psatd
from .Particles import newspecies, particles
from .WarpX import warpx
Expand Down
Loading
Loading