Coverage for pySDC/implementations/datatype_classes/cupy_mesh.py: 91%
66 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import cupy as cp
3try:
4 from mpi4py import MPI
5except ImportError:
6 MPI = None
8from pySDC.implementations.datatype_classes.mesh import MultiComponentMeshMixin
10try:
11 from pySDC.helpers.NCCL_communicator import NCCLComm
12except ImportError:
13 NCCLComm = None
16class cupy_mesh(cp.ndarray):
17 """
18 CuPy-based datatype for serial or parallel meshes.
19 """
21 comm = None
22 xp = cp
24 def __new__(cls, init, val=0.0, **kwargs):
25 """
26 Instantiates new datatype. This ensures that even when manipulating data, the result is still a mesh.
28 Args:
29 init: either another mesh or a tuple containing the dimensions, the communicator and the dtype
30 val: value to initialize
32 Returns:
33 obj of type mesh
35 """
36 if isinstance(init, cupy_mesh):
37 obj = cp.ndarray.__new__(cls, shape=init.shape, dtype=init.dtype, **kwargs)
38 obj[:] = init[:]
39 # a fresh array of its own, so `__array_finalize__` has nothing to inherit from and the
40 # communicator has to be carried over by hand
41 obj.comm = init.comm
42 elif (
43 isinstance(init, tuple)
44 and (init[1] is None or isinstance(init[1], MPI.Intracomm) or isinstance(init[1], NCCLComm))
45 and isinstance(init[2], cp.dtype)
46 ):
47 obj = cp.ndarray.__new__(cls, init[0], dtype=init[2], **kwargs)
48 obj.fill(val)
49 obj.comm = init[1]
50 else:
51 raise NotImplementedError(type(init))
52 return obj
54 def __array_ufunc__(self, ufunc, method, *inputs, out=None, **kwargs):
55 """
56 Overriding default ufunc, cf. https://numpy.org/doc/stable/user/basics.subclassing.html#array-ufunc-for-ufuncs
58 Only NumPy ufuncs come through here: CuPy's own `+=` calls its kernel with `out=self` directly and
59 is in place anyway. Outputs are handed on like the inputs, so `np.add(a, b, out=a)` writes into `a`.
60 """
61 args = [input_.view(cp.ndarray) if isinstance(input_, cupy_mesh) else input_ for input_ in inputs]
62 if out is not None:
63 kwargs['out'] = tuple(
64 [output.view(cp.ndarray) if isinstance(output, cupy_mesh) else output for output in out]
65 )
67 results = super(cupy_mesh, self).__array_ufunc__(ufunc, method, *args, **kwargs)
69 # written into what the caller gave, which keeps its own type and communicator
70 if out is not None:
71 return out[0] if len(out) == 1 else out
73 results = results.view(type(self))
75 # the inputs were viewed as plain arrays just above, so the result has no communicator to
76 # inherit through `__array_finalize__`; carry this one's over explicitly. A reduction or a
77 # comparison can return a scalar rather than an array, and a scalar takes no attributes.
78 if isinstance(results, cp.ndarray):
79 results.comm = self.comm
81 return results
83 def __array_finalize__(self, obj):
84 """
85 Carry the communicator onto every array derived from this one.
87 CuPy honours this hook, so a slice keeps the communicator. What reaches here from a ufunc
88 or a copy is a plain `cupy.ndarray` with nothing to carry, which is why those propagate it
89 themselves.
90 """
91 if obj is not None:
92 self.comm = getattr(obj, 'comm', None)
94 def __abs__(self):
95 """
96 Overloading the abs operator
98 Returns:
99 float: absolute maximum of all mesh values
100 """
101 # take absolute values of the mesh values
102 local_absval = cp.max(cp.ndarray.__abs__(self))
104 if self.comm is not None:
105 if self.comm.Get_size() > 1:
106 global_absval = local_absval * 0
107 if isinstance(self.comm, NCCLComm):
108 self.comm.Allreduce(sendbuf=local_absval, recvbuf=global_absval, op=MPI.MAX)
109 else:
110 global_absval = self.comm.allreduce(sendobj=float(local_absval), op=MPI.MAX)
111 else:
112 global_absval = local_absval
113 else:
114 global_absval = local_absval
116 return float(global_absval)
118 def copy(self, *args, **kwargs):
119 """
120 Copy the mesh, keeping its type.
122 CuPy's ``ndarray`` is a Cython extension type whose ``copy`` returns the base class rather than
123 ``type(self)``, unlike NumPy's. Without this, a copy of a mesh silently stops being one and loses the
124 components along with the ``__abs__`` that does the MPI reduction.
126 Returns:
127 mesh: copy of the mesh, of the same type
128 """
129 copied = super().copy(*args, **kwargs).view(type(self))
130 copied.comm = self.comm
131 return copied
133 def __deepcopy__(self, memo=None):
134 """
135 CuPy routes ``copy.deepcopy`` through ``copy``, so the type has to be kept here as well.
136 """
137 return self.copy()
139 def isend(self, dest=None, tag=None, comm=None):
140 """
141 Routine for sending data forward in time (non-blocking)
143 Args:
144 dest (int): target rank
145 tag (int): communication tag
146 comm: communicator
148 Returns:
149 request handle
150 """
151 return comm.Issend(self[:], dest=dest, tag=tag)
153 def irecv(self, source=None, tag=None, comm=None):
154 """
155 Routine for receiving in time
157 Args:
158 source (int): source rank
159 tag (int): communication tag
160 comm: communicator
162 Returns:
163 None
164 """
165 return comm.Irecv(self[:], source=source, tag=tag)
167 def bcast(self, root=None, comm=None):
168 """
169 Routine for broadcasting values
171 Args:
172 root (int): process with value to broadcast
173 comm: communicator
175 Returns:
176 broadcasted values
177 """
178 comm.Bcast(self[:], root=root)
179 return self
182class CuPyMultiComponentMesh(MultiComponentMeshMixin, cupy_mesh):
183 """CuPy-based mesh with multiple components, see ``MultiComponentMeshMixin``."""
186class imex_cupy_mesh(CuPyMultiComponentMesh):
187 """CuPy mesh with an implicit part ``impl`` and an explicit part ``expl``, for IMEX right-hand sides."""
189 components = ['impl', 'expl']
192class comp2_cupy_mesh(CuPyMultiComponentMesh):
193 """CuPy mesh with two parts ``comp1`` and ``comp2``, for multi-implicit right-hand sides."""
195 components = ['comp1', 'comp2']