diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 2b1a685a1..4bd5fd377 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -234,6 +234,34 @@ Conventions, as in struphy: - GPU tests are skipped without CuPy and a GPU (`cunumpy.kernel_testing.requires_cupy`). Until a GPU runner exists, they are run by hand on an H100 before a PR that touches CUDA code is merged, and the PR description says so. +## Inner products on the device + +`StencilVectorSpace.inner` (and so `StencilVector.inner`, `BlockVectorSpace.inner`, `dot_inner`) returns a 0-d +CuPy array for device data, in the serial and the MPI case (the `Allreduce` works on the device buffers, after +`synchronize_for_mpi`); before, the serial case copied the 8-byte result to the host. The result is a copy of the +reduction buffer, which the next inner product with the same vector overwrites. On NumPy it is a NumPy scalar, as +before. + +- `StencilVectorSpace.axpy` with a 0-d device array computes `y += a * x` with array operations (the axpy kernel + takes `alpha` by value, so calling it would copy `a` to the host and wait for the device). +- CG, PCG, BiCG, BiCGStab and PBiCGStab keep their step sizes on the device; the only copy per iteration is the + residual norm of the convergence test (`solvers._host`). BiCGStab tested the residual twice per iteration (once + redundantly); it now tests it once. MINRES, LSMR and the Uzawa solver do their scalar recurrences on the host + and copy every inner product, as before; GMRES keeps its Arnoldi inner products on the device. +- Copies to the host per iteration, counted under the fake CuPy (`linalg/tests/test_inner_on_device.py`, which + also counts implicit conversions such as `float()`, invisible to `count_transfers`): + + | Solver | before | after | + | --- | --- | --- | + | CG | 2 | 1 | + | PCG | 3 | 1 | + | BiCG | 4 | 1 | + | BiCGStab | 6 | 1 | + | PBiCGStab | 5 | 1 | + | MINRES, LSMR | 3 | 3 | + + Iteration counts and solutions are identical to NumPy. + ## Open questions - **Setup kernels on the GPU.** B-spline, field evaluation and DOF kernels still run on the host through @@ -245,8 +273,11 @@ Conventions, as in struphy: layout first. - **Complex data on the device.** Not needed by struphy so far; would need a second CUDA kernel per folder (or dtype dispatch in `cunumpy.kernels.Kernel`). -- **`inner` reduction.** One `atomicAdd` per block of 256 threads, then a copy of the 8-byte result to the host in - the serial case (in the parallel case it goes into the MPI reduction). To be measured on the H100. +- **`inner` reduction.** One `atomicAdd` per block of 256 threads; the result stays on the device (see + [Inner products on the device](#inner-products-on-the-device)). To be measured on the H100. +- **Fewer convergence tests.** The Krylov solvers still copy the residual norm to the host once per iteration. + Testing it every k iterations would remove most of these synchronizations, at the cost of up to k - 1 extra + iterations (which must not divide by a zero residual) and of iteration counts that differ from NumPy. - **Interface matrices** (`StencilInterfaceMatrix`) and the remaining stencil kernels (`stencil2coo`, ...) still use `PyccelKernel` with host copies. - **FFT on the device.** `DistributedFFT` and friends still stage their data through the host (SciPy FFT); they diff --git a/feectools/linalg/basic.py b/feectools/linalg/basic.py index 1fdcc4431..0ddf67a9c 100644 --- a/feectools/linalg/basic.py +++ b/feectools/linalg/basic.py @@ -93,9 +93,11 @@ def inner(self, x, y): Returns ------- - float | complex + float | complex | cupy.ndarray The scalar product of the two vectors. Note that inner(x, x) is a non-negative real number which is zero if and only if x = 0. + For vectors with device (CuPy) data, a 0-d device array: the + result stays on the device. """ diff --git a/feectools/linalg/block.py b/feectools/linalg/block.py index 526f2366d..541f1c4a0 100644 --- a/feectools/linalg/block.py +++ b/feectools/linalg/block.py @@ -112,9 +112,11 @@ def inner(self, x, y): Returns ------- - float | complex + float | complex | cupy.ndarray The scalar product of the two vectors. Note that inner(x, x) is a non-negative real number which is zero if and only if x = 0. + For vectors with device (CuPy) data, a 0-d device array: the + result stays on the device. """ @@ -134,7 +136,8 @@ def axpy(self, a, x, y): Parameters ---------- a : scalar - The scaling coefficient needed for the operation. + The scaling coefficient needed for the operation (a 0-d device + array is accepted, see `StencilVectorSpace.axpy`). x : BlockVector The vector which is not modified by this function. diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index 127e115a1..88328f575 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -25,6 +25,27 @@ 'GMRES', ) +#=============================================================================== +# Scalars on the device +# +# On the CuPy backend `Vector.inner` returns a 0-d device array. The Krylov +# solvers below (CG, PCG, BiCG, BiCGStab, PBiCGStab) keep their step sizes +# (alpha, beta, omega, ...) on the device: arithmetic with 0-d device arrays +# and `axpy` with a device scalar do not copy to the host. The only copy per +# iteration is the 8-byte residual norm of the convergence test, made explicit +# with `_host`. MINRES, LSMR and the Uzawa solver do their scalar recurrences +# on the host and copy every inner product (as before). On the NumPy backend +# nothing changes: both helpers return their argument as it is. +#=============================================================================== +def _host(s): + """A device scalar (0-d CuPy array) as a NumPy scalar (one copy to the host); anything else unchanged.""" + return xp.to_numpy(s)[()] if xp.is_gpu(s) else s + + +def _sqrt(s): + """Square root that keeps a device scalar on the device; `math.sqrt` for host scalars.""" + return xp.sqrt(s) if xp.is_gpu(s) else sqrt(s) + #=============================================================================== def inverse(A, solver, **kwargs): """ @@ -205,11 +226,12 @@ def solve(self, b, out=None): print( "+ Iter. # | L2-norm of residual |") print( "+---------+---------------------+") template = "| {:7d} | {:19.2e} |" - print(template.format(1, sqrt(am))) + print(template.format(1, sqrt(_host(am)))) # Iterate to convergence for m in range(2, maxiter+1): - if am < tol_sqr: + # the only copy to the host per iteration (with device vectors) + if _host(am) < tol_sqr: m -= 1 break A.dot(p, out=v) @@ -223,12 +245,13 @@ def solve(self, b, out=None): p += r am = am1 if verbose: - print(template.format(m, sqrt(am))) + print(template.format(m, sqrt(_host(am)))) if verbose: print( "+---------+---------------------+") # Convergence information + am = _host(am) self._info = {'niter': m, 'success': am < tol_sqr, 'res_norm': sqrt(am) } if recycle: @@ -368,12 +391,13 @@ def solve(self, b, out=None): print( "+ Iter. # | L2-norm of residual |") print( "+---------+---------------------+") template = "| {:7d} | {:19.2e} |" - print( template.format(1, sqrt(nrmr_sqr))) + print( template.format(1, sqrt(_host(nrmr_sqr)))) # Iterate to convergence for k in range(2, maxiter+1): - if nrmr_sqr < tol_sqr: + # the only copy to the host per iteration (with device vectors) + if _host(nrmr_sqr) < tol_sqr: k -= 1 break @@ -395,12 +419,13 @@ def solve(self, b, out=None): am = am1 if verbose: - print( template.format(k, sqrt(nrmr_sqr))) + print( template.format(k, sqrt(_host(nrmr_sqr)))) if verbose: print( "+---------+---------------------+") # Convergence information + nrmr_sqr = _host(nrmr_sqr) self._info = {'niter': k, 'success': nrmr_sqr < tol_sqr, 'res_norm': sqrt(nrmr_sqr) } if recycle: @@ -540,7 +565,8 @@ def solve(self, b, out=None): # Iterate to convergence for m in range(1, maxiter + 1): - if res_sqr < tol_sqr: + # the only copy to the host per iteration (with device vectors) + if _host(res_sqr) < tol_sqr: m -= 1 break @@ -585,12 +611,13 @@ def solve(self, b, out=None): ps += rs if verbose: - print( template.format(m, sqrt(res_sqr)) ) + print( template.format(m, sqrt(_host(res_sqr))) ) if verbose: print( "+---------+---------------------+") # Convergence information + res_sqr = _host(res_sqr) self._info = {'niter': m, 'success': res_sqr < tol_sqr, 'res_norm': sqrt(res_sqr)} if recycle: @@ -729,12 +756,12 @@ def solve(self, b, out=None): print("+---------+---------------------+") template = "| {:7d} | {:19.2e} |" - # Iterate to convergence - for m in range(1, maxiter + 1): - - if res_sqr < tol_sqr: - m -= 1 - break + # Iterate to convergence. The residual is tested once per iteration, + # after its update (the only copy to the host per iteration with device + # vectors); before the first iteration it is tested here. + n_iter = 0 if _host(res_sqr) < tol_sqr else maxiter + m = 0 + for m in range(1, n_iter + 1): # ----------------------- # MATRIX-VECTOR PRODUCTS @@ -771,7 +798,7 @@ def solve(self, b, out=None): # ||r||_2 := (r, r) res_sqr = r.inner(r).real - if res_sqr < tol_sqr: + if _host(res_sqr) < tol_sqr: break # b := a / w * (r0, r)_{m+1} / (r0, r)_m @@ -783,12 +810,13 @@ def solve(self, b, out=None): p.mul_iadd(-b * w, v) if verbose: - print(template.format(m, sqrt(res_sqr))) + print(template.format(m, sqrt(_host(res_sqr)))) if verbose: print("+---------+---------------------+") # Convergence information + res_sqr = _host(res_sqr) self._info = {'niter': m, 'success': res_sqr < tol_sqr, 'res_norm': sqrt(res_sqr)} if recycle: @@ -963,7 +991,8 @@ def solve(self, b, out=None): # iterate to convergence or maximum number of iterations niter = 0 - while res_sqr > tol_sqr and niter < maxiter: + # the only copy to the host per iteration (with device vectors) + while _host(res_sqr) > tol_sqr and niter < maxiter: # v = A @ pp, vp = PC @ v, alphap = rhop/(vp.rp0) A.dot(pp, out=v) @@ -1016,12 +1045,13 @@ def solve(self, b, out=None): niter += 1 if verbose: - print(template.format(niter, sqrt(res_sqr))) + print(template.format(niter, sqrt(_host(res_sqr)))) if verbose: print("+---------+---------------------+") # convergence information + res_sqr = _host(res_sqr) self._info = {'niter': niter, 'success': res_sqr < tol_sqr, 'res_norm': sqrt(res_sqr)} @@ -1171,7 +1201,7 @@ def solve(self, b, out=None): y *= -1.0 y.copy(out=res_old) # res = b - A*x - beta = sqrt(res_old.inner(res_old)) + beta = sqrt(_host(res_old.inner(res_old))) # Initialize other quantities oldb = 0 @@ -1215,7 +1245,7 @@ def solve(self, b, out=None): if itn >= 2: y.mul_iadd(-(beta/oldb), res_old) - alfa = v.inner(y) + alfa = _host(v.inner(y)) y.mul_iadd(-(alfa/beta), res_new) # We put res_new in res_old and y in res_new @@ -1223,7 +1253,7 @@ def solve(self, b, out=None): y.copy(out=res_new) oldb = beta - beta = sqrt(res_new.inner(res_new)) + beta = sqrt(_host(res_new.inner(res_new))) tnorm2 += alfa**2 + oldb**2 + beta**2 # Apply previous rotation Qk-1 to get @@ -1270,7 +1300,7 @@ def solve(self, b, out=None): # Estimate various norms and test for convergence. Anorm = sqrt(tnorm2) - ynorm = sqrt(x.inner(x)) + ynorm = sqrt(_host(x.inner(x))) rnorm = phibar if ynorm == 0 or Anorm == 0:test1 = inf @@ -1486,16 +1516,16 @@ def solve(self, b, out=None): btol = tol b.copy(out=u) - normb = sqrt(b.inner(b).real) + normb = sqrt(_host(b.inner(b)).real) A.dot(x, out=u_work) u -= u_work - beta = sqrt(u.inner(u).real) + beta = sqrt(_host(u.inner(u)).real) if beta > 0: u *= (1 / beta) At.dot(u, out=v) - alpha = sqrt(v.inner(v).real) + alpha = sqrt(_host(v.inner(v)).real) else: x.copy(out=v) alpha = 0 @@ -1558,14 +1588,14 @@ def solve(self, b, out=None): u *= -alpha A.dot(v, out=u_work) u += u_work - beta = sqrt(u.inner(u).real) + beta = sqrt(_host(u.inner(u)).real) if beta > 0: u *= (1 / beta) v *= -beta At.dot(u, out=v_work) v += v_work - alpha = sqrt(v.inner(v).real) + alpha = sqrt(_host(v.inner(v)).real) if alpha > 0:v *= (1 / alpha) # At this point, beta = beta_{k+1}, alpha = alpha_{k+1}. @@ -1642,7 +1672,7 @@ def solve(self, b, out=None): # Compute norms for convergence testing. normar = abs(zetabar) - normx = sqrt(x.inner(x).real) + normx = sqrt(_host(x.inner(x)).real) # Now use these norms to estimate certain other quantities, # some of which will be small near a solution. @@ -1807,7 +1837,7 @@ def solve(self, b, out=None): A.dot( x , out=r) r -= b - am = sqrt(r.inner(r).real) + am = sqrt(_host(r.inner(r).real)) if am < tol: self._info = {'niter': 1, 'success': am < tol, 'res_norm': am } return x @@ -1884,7 +1914,7 @@ def arnoldi(self, k, p): h[i] = p.inner(self._Q[i]) p.mul_iadd(-h[i], self._Q[i]) - h[k+1] = sqrt(p.inner(p).real) + h[k+1] = _sqrt(p.inner(p).real) p /= h[k+1] # Normalize vector if len(self._Q) > k + 1: @@ -2007,7 +2037,7 @@ def solve(self, b, out=None): # constraint residual: R = B1*u + B2*ue - g R = B1.dot(u) + B2.dot(ue) - g - residual_norm = sqrt(R.inner(R).real) + residual_norm = sqrt(_host(R.inner(R)).real) if verbose: print(template.format(iteration, residual_norm)) @@ -2017,7 +2047,7 @@ def solve(self, b, out=None): # pressure update: steepest descent step size S_R = B1.dot(A11inv.dot(B1.T.dot(R))) + B2.dot(A22inv.dot(B2.T.dot(R))) - alpha = R.inner(R).real / R.inner(S_R).real + alpha = _host(R.inner(R)).real / _host(R.inner(S_R)).real p += alpha * R if verbose: diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index ee7092fc9..cc8299189 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -300,9 +300,12 @@ def inner(self, x, y): Returns ------- - float | complex + float | complex | cupy.ndarray The scalar product of the two vectors. Note that inner(x, x) is a non-negative real number which is zero if and only if x = 0. + A NumPy scalar for NumPy data; for device (CuPy) data a 0-d + device array, so that the result stays on the device (no copy + to the host, no synchronization). """ @@ -324,10 +327,14 @@ def inner(self, x, y): self.cart.global_comm.Allreduce((x._dot_send_data, self.mpi_type), (x._dot_recv_data, self.mpi_type), op=MPI.SUM ) - return x._dot_recv_data[0] - else: - # a NumPy scalar on both backends - return xp.to_numpy(res)[0] + res = x._dot_recv_data + + if xp.is_gpu(res): + # A 0-d device array. Copied (on the device), because res is a buffer + # of x that the next inner product with x overwrites. + return res[0].copy() + # a NumPy scalar + return res[0] # ... def axpy(self, a, x, y): @@ -339,7 +346,9 @@ def axpy(self, a, x, y): Parameters ---------- a : scalar - The scaling coefficient needed for the operation. + The scaling coefficient needed for the operation. A Python or + NumPy scalar, or a 0-d device array (e.g. the result of `inner` + on device data), which is not copied to the host. x : StencilVector The vector which is not modified by this function. @@ -352,6 +361,18 @@ def axpy(self, a, x, y): assert x._space is self assert y._space is self + if xp.is_gpu(a): + # A device scalar, e.g. the result of an inner product: converting it + # to a Python scalar for the kernel would copy it to the host and wait + # for the device, so y += a * x is computed with array operations. + if self.dtype != complex and xp.iscomplexobj(a): + raise TypeError('A complex scalar was given in a real case') + y._data += a * x._data + for axis, ext in self.interfaces: + y._interface_data[axis, ext] += a * x._interface_data[axis, ext] + y._sync = x._sync and y._sync + return + if self.dtype == complex: a = complex(a) else: diff --git a/feectools/linalg/tests/inner_on_device_child.py b/feectools/linalg/tests/inner_on_device_child.py new file mode 100644 index 000000000..6745b730a --- /dev/null +++ b/feectools/linalg/tests/inner_on_device_child.py @@ -0,0 +1,149 @@ +"""Child process of `test_inner_on_device.py`: inner products and Krylov solvers under cunumpy's fake CuPy. + +Run with ``CUNUMPY_FAKE_CUPY=1 CUNUMPY_BACKEND=cupy MAYBEMPI=0`` in a clean environment (the fake must be +installed before cunumpy picks its backend, which is why this runs in its own process). Prints one JSON object. + +The fake CuPy cannot launch CUDA kernels, so `Kernel` calls with device arrays run the kernel's host +implementation on the host buffers of the fake device arrays (no copy, like a launch on the device data). +Two counts are reported: ``to_host`` of `cunumpy.profiling.count_transfers` (copies made through cunumpy, +e.g. the convergence test of a solver) and all device -> host copies, counted by wrapping the methods of the +fake ``cupy.ndarray`` that copy (``.get()``, which `to_numpy` uses, ``.item()``, ``float()``, ``bool()``, ...): +`count_transfers` does not see the implicit conversions. +""" +import json +import sys + +import cunumpy as xp +import numpy as np +from cunumpy import _fake_cupy +from cunumpy._dispatch import Kernel + +assert _fake_cupy.is_active() and xp.get_backend() == "cupy", "run with CUNUMPY_FAKE_CUPY=1 CUNUMPY_BACKEND=cupy" + +import cupy # noqa: E402 (the fake) + + +# --- device kernel launches: the host implementation on the device buffers -------------------------------- +def _buffer(a): + return a._a if isinstance(a, cupy.ndarray) else a + + +def _launch(self, *args, **launch): + if any(isinstance(a, cupy.ndarray) for a in args): + return self.host_kernel.kernel(*[_buffer(a) for a in args]) + return self.host_kernel.kernel(*args) + + +Kernel.__call__ = _launch + +# --- every device -> host copy, implicit or not --------------------------------------------------------------- +COPIES = [] + + +def _counting(name, method): + def wrapper(self, *args, **kwargs): + COPIES.append(name) + return method(self, *args, **kwargs) + return wrapper + + +for _name in ("__bool__", "__float__", "__complex__", "__int__", "__index__", "__format__", "get"): + setattr(cupy.ndarray, _name, _counting(_name, getattr(cupy.ndarray, _name))) +cupy.ndarray.item = _counting("item", lambda self, *args: self._a.item(*args)) + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition # noqa: E402 +from feectools.linalg.block import BlockVector, BlockVectorSpace # noqa: E402 +from feectools.linalg.solvers import inverse # noqa: E402 +from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace # noqa: E402 + +N, P = 40, 2 +TOL = 1e-12 + + +def problem(): + """A symmetric positive definite banded matrix, an exact solution and the right-hand side, on the active backend.""" + D = DomainDecomposition([N - P], [False]) + C = CartDecomposition(D, [N], [np.array([0])], [np.array([N - 1])], [P], [1]) + V = StencilVectorSpace(C, dtype=float) + A = StencilMatrix(V, V) + A[:, -P:0] = -1 + A[:, 0:1] = 2 * P + A[:, 1:P + 1] = -1 + A.remove_spurious_entries() + xe = StencilVector(V) + xe[0:N] = xp.asarray(np.random.default_rng(0).random(N)) + return V, A, xe, A.dot(xe) + + +def solver_kwargs(name, A): + return {"pc": A.diagonal(inverse=True)} if name in ("pcg", "pbicgstab") else {} + + +def host(a): + return xp.to_numpy(a).tolist() + + +def check_inner(): + V, A, xe, b = problem() + y = b.copy() + out = {} + s1 = xe.inner(y) + s2 = xe.inner(xe) # reuses the reduction buffer of xe + out["stencil_type"] = type(s1).__module__ + "." + type(s1).__name__ + out["stencil_ndim"] = s1.ndim + out["independent"] = host(s1) != host(s2) + W = BlockVectorSpace(V, V) + u = BlockVector(W, blocks=[xe, y]) + sb = u.inner(u) + out["block_type"] = type(sb).__module__ + "." + type(sb).__name__ + out["block_value"] = host(sb) + # axpy with a device scalar, no copy to the host + z = y.copy() + COPIES.clear() + with xp.profiling.count_transfers() as c: + z.mul_iadd(s1, xe) + out["axpy_transfers"] = c.to_host + len(COPIES) + out["axpy_value"] = host(z._data) + out["inner_values"] = [host(s1), host(s2)] + return out + + +def run_solvers(names): + out = {} + for name in names: + V, A, xe, b = problem() + kw = solver_kwargs(name, A) + inverse(A, name, tol=TOL, maxiter=1000, **kw).solve(b) # warm up + counts = {} + for maxiter in (5, 15): + S = inverse(A, name, tol=1e-30, maxiter=maxiter, **kw) + COPIES.clear() + with xp.profiling.count_transfers() as c: + S.solve(b) + counts[maxiter] = (c.to_host, len(COPIES), c.to_device) + S = inverse(A, name, tol=TOL, maxiter=1000, **kw) + x = S.solve(b) + info = S.get_info() + out[name] = { + # per iteration: the difference between 15 and 5 iterations, over 10 iterations + "to_host_per_iter": (counts[15][0] - counts[5][0]) / 10, + "copies_per_iter": (counts[15][1] - counts[5][1]) / 10, + "to_device": counts[15][2], + "x": host(x._data), + "niter": int(info["niter"]), + "res_norm": float(info["res_norm"]), + "info_types": sorted({type(v).__module__.split(".")[0] for v in info.values()}), + } + # the NumPy reference, same process + with xp.use_backend("numpy"): + V, A, xe, b = problem() + S = inverse(A, name, tol=TOL, maxiter=1000, **solver_kwargs(name, A)) + x = S.solve(b) + out[name]["numpy_x"] = x._data.tolist() + out[name]["numpy_niter"] = int(S.get_info()["niter"]) + return out + + +if __name__ == "__main__": + result = {"inner": check_inner(), "solvers": run_solvers(sys.argv[1:])} + print("RESULT " + json.dumps(result)) diff --git a/feectools/linalg/tests/test_inner_on_device.py b/feectools/linalg/tests/test_inner_on_device.py new file mode 100644 index 000000000..451df1906 --- /dev/null +++ b/feectools/linalg/tests/test_inner_on_device.py @@ -0,0 +1,160 @@ +"""Inner products stay on the device, and the Krylov solvers keep their scalars there. + +On the CuPy backend `StencilVectorSpace.inner` (and `BlockVectorSpace.inner`) return a 0-d device array, `axpy` +accepts it without copying it to the host, and CG, PCG, BiCG, BiCGStab and PBiCGStab copy one scalar per +iteration to the host (the residual norm of the convergence test). On the NumPy backend nothing changes. + +Without a GPU the CuPy path runs under cunumpy's fake CuPy, in a child process +(`inner_on_device_child.py`); the GPU tests are skipped. +""" +import json +import os +import subprocess +import sys +from pathlib import Path + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.linalg.block import BlockVector, BlockVectorSpace +from feectools.linalg.solvers import inverse +from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace + +CHILD = Path(__file__).with_name("inner_on_device_child.py") +DEVICE_SOLVERS = ("cg", "pcg", "bicg", "bicgstab", "pbicgstab") +HOST_SOLVERS = ("minres", "lsmr", "gmres") + + +def problem(n=40, p=2): + D = DomainDecomposition([n - p], [False]) + C = CartDecomposition(D, [n], [np.array([0])], [np.array([n - 1])], [p], [1]) + V = StencilVectorSpace(C, dtype=float) + A = StencilMatrix(V, V) + A[:, -p:0] = -1 + A[:, 0:1] = 2 * p + A[:, 1:p + 1] = -1 + A.remove_spurious_entries() + xe = StencilVector(V) + xe[0:n] = xp.asarray(np.random.default_rng(0).random(n)) + return V, A, xe, A.dot(xe) + + +# =============================================================================== +# NumPy backend: unchanged +# =============================================================================== +def test_inner_numpy_returns_numpy_scalar(): + with xp.use_backend("numpy"): + V, A, xe, b = problem() + s1 = xe.inner(b) + s2 = xe.inner(xe) # reuses the reduction buffer of xe + assert isinstance(s1, np.floating) and isinstance(s2, np.floating) + assert s1 == pytest.approx(float(np.dot(xe.toarray(), b.toarray())), rel=1e-14) + assert s2 == pytest.approx(float(np.dot(xe.toarray(), xe.toarray())), rel=1e-14) + + W = BlockVectorSpace(V, V) + u = BlockVector(W, blocks=[xe, b]) + assert isinstance(u.inner(u), np.floating) + + +@pytest.mark.parametrize("solver", DEVICE_SOLVERS + HOST_SOLVERS) +def test_solver_info_numpy(solver): + """The convergence information is made of host scalars, as before.""" + with xp.use_backend("numpy"): + V, A, xe, b = problem() + kw = {"pc": A.diagonal(inverse=True)} if solver in ("pcg", "pbicgstab") else {} + S = inverse(A, solver, tol=1e-12, maxiter=1000, **kw) + x = S.solve(b) + info = S.get_info() + assert info["success"] + assert isinstance(info["res_norm"], float) + assert np.allclose(x.toarray(), xe.toarray(), atol=1e-10) + + +# =============================================================================== +# CuPy backend without a GPU: fake CuPy in a child process +# =============================================================================== +def fake_cupy_env(): + """A clean environment for a serial child process with the fake CuPy (no MPI launcher variables).""" + prefixes = ("OMPI_", "PMIX_", "PMI_", "HYDRA_", "I_MPI_", "SLURM_", "CUNUMPY_") + env = {k: v for k, v in os.environ.items() if not k.startswith(prefixes)} + env.update(CUNUMPY_FAKE_CUPY="1", CUNUMPY_BACKEND="cupy", MAYBEMPI="0") + root = str(Path(__file__).resolve().parents[3]) + env["PYTHONPATH"] = os.pathsep.join([root, env.get("PYTHONPATH", "")]).rstrip(os.pathsep) + return env + + +@pytest.fixture(scope="module") +def fake_cupy_result(): + try: + import cupy # noqa: F401 + except ImportError: + pass + else: + pytest.skip("CuPy is installed: the fake CuPy cannot be used") + proc = subprocess.run( + [sys.executable, str(CHILD), *DEVICE_SOLVERS, *HOST_SOLVERS], + env=fake_cupy_env(), capture_output=True, text=True, timeout=600, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + line = [ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT ")][-1] + return json.loads(line[len("RESULT "):]) + + +def test_inner_returns_device_scalar(fake_cupy_result): + r = fake_cupy_result["inner"] + assert r["stencil_type"] == "cupy.ndarray" and r["stencil_ndim"] == 0 + assert r["block_type"] == "cupy.ndarray" + # the result is a copy, not a view of the reduction buffer that the next inner product overwrites + assert r["independent"] + + +def test_axpy_with_device_scalar_does_not_copy(fake_cupy_result): + assert fake_cupy_result["inner"]["axpy_transfers"] == 0 + + +@pytest.mark.parametrize("solver", DEVICE_SOLVERS) +def test_one_host_copy_per_iteration(fake_cupy_result, solver): + """Before: 2 (CG), 3 (PCG), 4 (BiCG), 6 (BiCGStab), 5 (PBiCGStab) copies per iteration.""" + r = fake_cupy_result["solvers"][solver] + assert r["to_host_per_iter"] == 1 + assert r["copies_per_iter"] == 1 + assert r["to_device"] == 0 + assert r["info_types"] in (["builtins"], ["builtins", "numpy"]) + + +@pytest.mark.parametrize("solver", DEVICE_SOLVERS + HOST_SOLVERS) +def test_same_result_as_numpy(fake_cupy_result, solver): + r = fake_cupy_result["solvers"][solver] + assert r["niter"] == r["numpy_niter"] + np.testing.assert_allclose(r["x"], r["numpy_x"], rtol=1e-12, atol=1e-14) + + +# =============================================================================== +# GPU +# =============================================================================== +@requires_cupy +def test_inner_on_gpu(): + with xp.use_backend("cupy"): + V, A, xe, b = problem() + s = xe.inner(b) + assert xp.is_gpu(s) and s.ndim == 0 + assert float(s) == pytest.approx(float(xp.to_numpy(xe._data) @ xp.to_numpy(b._data)), rel=1e-12) + + +@requires_cupy +@pytest.mark.parametrize("solver", DEVICE_SOLVERS) +def test_one_host_copy_per_iteration_on_gpu(solver): + with xp.use_backend("cupy"): + V, A, xe, b = problem() + kw = {"pc": A.diagonal(inverse=True)} if solver in ("pcg", "pbicgstab") else {} + inverse(A, solver, tol=1e-12, maxiter=1000, **kw).solve(b) + counts = [] + for maxiter in (5, 15): + S = inverse(A, solver, tol=1e-30, maxiter=maxiter, **kw) + with xp.profiling.count_transfers() as c: + S.solve(b) + counts.append(c.to_host) + assert (counts[1] - counts[0]) == 10 diff --git a/feectools/linalg/tests/test_mpi_device.py b/feectools/linalg/tests/test_mpi_device.py index 70768c213..a9eec09be 100644 --- a/feectools/linalg/tests/test_mpi_device.py +++ b/feectools/linalg/tests/test_mpi_device.py @@ -282,6 +282,31 @@ def test_fewer_diagonals_dot_and_transpose_match_global_reference(): assert abs(lhs - rhs) <= 1e-12 * abs(rhs) +# =============================================================================== +def test_inner_result_is_a_copy_of_the_reduction_buffer(): + """The result of the MPI reduction survives the next inner product with the same vector. + + `inner` reduces into a buffer of the first vector; on the device it returns a 0-d + device array (no copy to the host), which must not be a view of that buffer. + """ + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + other = np.flipud(glob).copy() + + x = scatter(V, glob) + y = scatter(V, other) + + xy = x.inner(y) + xx = x.inner(x) + if xp.is_gpu(x._data): + assert xp.is_gpu(xy) and xy.ndim == 0 + else: + assert isinstance(xy, np.floating) + assert abs(float(xy) - float((glob * other).sum())) <= 1e-9 + assert abs(float(xx) - float((glob * glob).sum())) <= 1e-9 + + # =============================================================================== if __name__ == "__main__": import sys