batched CG on CPU?

[lucascolley]

Hi! Is gko::batch::solver::Cg batched even on CPU? I see the CUDA kernel is at https://github.com/ginkgo-project/ginkgo/blob/develop/common/cuda_hip/solver/batch_cg_kernels.hpp, but this won't be used on a CPU-only machine, right? If it is batched, where is the source code for the algorithm equivalent to what the CUDA kernel implements?

[pratikvn]

Hi @lucascolley , the batched solvers on CPU backends are run either with gko::ReferenceExecutor or gko::OmpExecutor. The kernel for the OpenMP backend, for example is available here : https://github.com/ginkgo-project/ginkgo/blob/develop/omp/solver/batch_cg_kernels.cpp.

With the batched methods, the low-level BLAS operations, and the per-linear system CG solver are shared between reference and OMP backends and implemented here: https://github.com/ginkgo-project/ginkgo/blob/develop/reference/solver/batch_cg_kernels.hpp

As a user, you can solve as many systems in parallel as the number of OpenMP threads you spawn, and each system solve converges (terminates) independently, so you are encouraged to oversubscribe to maximize the occupancy.

[lucascolley]

thanks, makes sense! I think that answers all of my questions.

For some context (in case you are interested), I am exploring naïve Python-level batching of the iterative solvers available in SciPy as part of a master's project. For the most naïve form, the following diff seems pretty much sufficient:

diff --git a/scipy/sparse/linalg/_isolve/iterative.py b/scipy/sparse/linalg/_isolve/iterative.py
index 637f2d021f..b4084e30d4 100644
--- a/scipy/sparse/linalg/_isolve/iterative.py
+++ b/scipy/sparse/linalg/_isolve/iterative.py
@@ -399,13 +401,23 @@ def cg(A, b, x0=None, *, rtol=1e-5, atol=0., maxiter=None, M=None, callback=None
     rho_prev, p = None, None

     for iteration in range(maxiter):
-        if np.linalg.norm(r) < atol:  # Are we done?
+        converged = np.linalg.norm(r, axis=-1) < atol
+        if np.all(converged):
             return x, 0

         z = psolve(r)
         rho_cur = dotprod(r, z)
+
         if iteration > 0:
-            beta = rho_cur / rho_prev
+            beta = xpx.apply_where(
+                ~converged,
+                (rho_cur, rho_prev),
+                lambda cur, prev: cur / prev,
+                fill_value=0.0,
+                xp=np
+            )
+            beta = np.expand_dims(beta, axis=-1)
+
             p *= beta
             p += z
         else:  # First spin
@@ -413,7 +425,17 @@ def cg(A, b, x0=None, *, rtol=1e-5, atol=0., maxiter=None, M=None, callback=None
             p[:] = z[:]

         q = matvec(p)
-        alpha = rho_cur / dotprod(p, q)
+        c = dotprod(p, q)
+
+        alpha = xpx.apply_where(
+            ~converged,
+            (rho_cur, c),
+            lambda rc, c: rc / c,
+            fill_value=0.0,
+            xp=np
+        )
+        alpha = np.expand_dims(alpha, axis=-1)
+
         x += alpha*p
         r -= alpha*q
         rho_prev = rho_cur

However, there is the problem of wasted work for already converged systems, in the matvec application, but also in the convergence check and dot products.

I was wondering whether there was any prior art regarding 'zeroing-out' of already converged systems for a batched solver which keeps the steps of each system 'in-sync' like this, to avoid the redundant computation. However, IIUC, it sounds like the implementations in ginkgo keep each system in separate threads (or whatever the equivalent is for CUDA etc.) rather than performing the iterations in-sync.

[pratikvn]

In general, in addition to the wasted work, IIRC, we also faced stability issues, as over-iterating after some systems have can lead to divergence/blow-up, so numerically, it makes sense to exit after the system converges. For example, for our monolithic solvers (non-batched), where we can solve multiple rhs, it is necessary for us to "toggle-off" iterations for the rhs that have already converged.

For the batched case, we therefore made a deliberate choice to map one linear system solve to one work unit (OMP thread, CUDA thread block, SYCL workgroup etc), thereby each system converges independently, but dont waste resources (they are not left idle, and still do useful work) because we always have new systems to solve on the resources because we oversubscribe the work unit by throwing a large numer of linear systems, ideally much larger than the number work units, which also tends to hide any latencies.

[lucascolley]

right, makes sense. In the case of CG, I think that stability issues should be addressed by lines like this, which ensure that the alpha and beta steps are zero for converged systems:

beta = xpx.apply_where(
    ~converged,
    (rho_cur, rho_prev),
    lambda cur, prev: cur / prev,
    fill_value=0.0,
    xp=np
)

But I wonder whether there is no good solution to the wasted work problem, if the bottleneck is memory access rather than arithmetic.

Still, I suppose this approach would be valuable in the case that all systems take roughly the same amount of time to converge. I will have to investigate strategies for parallelism.

[pratikvn]

Yes, the optimal strategy might depend on the hardware you run on. On GPUs, you pay a high cost for parallel resource divergence, and for non-uniform memory access, so zeroing might might actually be the better strategy. On modern CPUs, it might be similar. But at the same time if you do some sort of a task based approach, wasting resources when some other work could have been performed, might not be optimal.

[lucascolley]

thanks again for your help!