ell_relax -> ell_grid_solve

Generalize the API to allow for multiple solver types. This is done in
preparation for the exact linear system inversion solver.

5 files changed, 166 insertions, 103 deletions

diff --git a/ell_relax.c b/ell_grid_solve.c
index 4ee2747..191fc8f 100644
--- a/ell_relax.c
+++ b/ell_grid_solve.c
@@ -26,7 +26,7 @@
 
 #include "common.h"
 #include "cpu.h"
-#include "ell_relax.h"
+#include "ell_grid_solve.h"
 #include "log.h"
 #include "mg2d_boundary.h"
 #include "mg2d_boundary_internal.h"
@@ -37,7 +37,11 @@ static const double relax_factors[FD_STENCIL_MAX] = {
     [1] = 1.0 / 7,
 };
 
-struct EllRelaxInternal {
+typedef struct EGSRelaxInternal {
+    double relax_factor;
+} EGSRelaxInternal;
+
+struct EGSInternal {
     ptrdiff_t stride;
     double *u_base;
     double *rhs_base;
@@ -55,7 +59,10 @@ struct EllRelaxInternal {
     ptrdiff_t residual_calc_offset;
     size_t residual_calc_size[2];
     double fd_factors[MG2D_DIFF_COEFF_NB];
-    double relax_factor;
+
+    union {
+        EGSRelaxInternal r;
+    };
 
     TPContext *tp_internal;
 };
@@ -222,8 +229,8 @@ static void residual_calc_line_s2_c(size_t linesize, double *dst, double *dst_ma
 
 static void residual_calc_task(void *arg, unsigned int job_idx, unsigned int thread_idx)
 {
-    EllRelaxContext  *ctx = arg;
-    EllRelaxInternal *priv = ctx->priv;
+    EGSContext  *ctx = arg;
+    EGSInternal *priv = ctx->priv;
     const ptrdiff_t offset = priv->residual_calc_offset + job_idx * priv->stride;
     const double *diff_coeffs[MG2D_DIFF_COEFF_NB];
 
@@ -236,9 +243,9 @@ static void residual_calc_task(void *arg, unsigned int job_idx, unsigned int thr
                              diff_coeffs, priv->fd_factors);
 }
 
-static void residual_calc(EllRelaxContext *ctx)
+static void residual_calc(EGSContext *ctx)
 {
-    EllRelaxInternal *priv = ctx->priv;
+    EGSInternal *priv = ctx->priv;
     double res_max = 0.0;
     int64_t start;
 
@@ -280,9 +287,9 @@ static void boundaries_apply_fixdiff(double *dst, const ptrdiff_t dst_stride[2],
     }
 }
 
-static void boundaries_apply(EllRelaxContext *ctx)
+static void boundaries_apply(EGSContext *ctx)
 {
-    EllRelaxInternal *priv = ctx->priv;
+    EGSInternal *priv = ctx->priv;
     const ptrdiff_t strides[2] = { 1, priv->stride };
     int64_t start;
 
@@ -385,39 +392,47 @@ static void boundaries_apply(EllRelaxContext *ctx)
 
 static void residual_add_task(void *arg, unsigned int job_idx, unsigned int thread_idx)
 {
-    EllRelaxContext   *ctx = arg;
-    EllRelaxInternal *priv = ctx->priv;
+    EGSContext   *ctx = arg;
+    EGSInternal *priv = ctx->priv;
     ptrdiff_t offset = job_idx * priv->stride;
 
     for (int idx0 = 0; idx0 < ctx->domain_size[0]; idx0++) {
         ptrdiff_t idx = job_idx * ctx->u_stride + idx0;
 
-        ctx->u[idx] += priv->relax_factor * ctx->residual[idx];
+        ctx->u[idx] += priv->r.relax_factor * ctx->residual[idx];
     }
 }
 
-int mg2di_ell_relax_step(EllRelaxContext *ctx)
+static int solve_relax_step(EGSContext *ctx)
 {
-    EllRelaxInternal *priv = ctx->priv;
-    const double cfl_fac = priv->relax_factor;
+    EGSRelaxContext *r = ctx->solver_data;
+    EGSInternal *priv = ctx->priv;
     int64_t start;
 
     start = gettime();
 
     tp_execute(ctx->tp, ctx->domain_size[1], residual_add_task, ctx);
 
-    ctx->time_correct += gettime() - start;
-    ctx->count_correct++;
+    r->time_correct += gettime() - start;
+    r->count_correct++;
 
     boundaries_apply(ctx);
     residual_calc(ctx);
 
     return 0;
 }
+int mg2di_egs_solve(EGSContext *ctx)
+{
+    switch (ctx->solver_type) {
+    case EGS_SOLVER_RELAXATION: return solve_relax_step(ctx);
+    }
 
-int mg2di_ell_relax_init(EllRelaxContext *ctx)
+    return -EINVAL;
+}
+
+int mg2di_egs_init(EGSContext *ctx)
 {
-    EllRelaxInternal *priv = ctx->priv;
+    EGSInternal *priv = ctx->priv;
     double *tmp;
     int ret;
 
@@ -452,11 +467,14 @@ int mg2di_ell_relax_init(EllRelaxContext *ctx)
         return -EINVAL;
     }
 
-    if (ctx->relax_factor == 0.0)
-        priv->relax_factor = relax_factors[ctx->fd_stencil - 1];
-    else
-        priv->relax_factor = ctx->relax_factor;
-    priv->relax_factor *= ctx->step[0] * ctx->step[0];
+    if (ctx->solver_type == EGS_SOLVER_RELAXATION) {
+        EGSRelaxContext *r = ctx->solver_data;
+        if (r->relax_factor == 0.0)
+            priv->r.relax_factor = relax_factors[ctx->fd_stencil - 1];
+        else
+            priv->r.relax_factor = r->relax_factor;
+        priv->r.relax_factor *= ctx->step[0] * ctx->step[0];
+    }
 
     priv->fd_factors[MG2D_DIFF_COEFF_00] = 1.0 /  fd_denoms[ctx->fd_stencil - 1][MG2D_DIFF_COEFF_00];
     priv->fd_factors[MG2D_DIFF_COEFF_10] = 1.0 / (fd_denoms[ctx->fd_stencil - 1][MG2D_DIFF_COEFF_10] * ctx->step[0]);
@@ -508,9 +526,9 @@ int mg2di_ell_relax_init(EllRelaxContext *ctx)
     return 0;
 }
 
-static int ell_relax_arrays_alloc(EllRelaxContext *ctx, const size_t domain_size[2])
+static int arrays_alloc(EGSContext *ctx, const size_t domain_size[2])
 {
-    EllRelaxInternal *priv = ctx->priv;
+    EGSInternal *priv = ctx->priv;
 
     const size_t ghosts = FD_STENCIL_MAX;
     const size_t size_padded[2] = {
@@ -561,10 +579,10 @@ static int ell_relax_arrays_alloc(EllRelaxContext *ctx, const size_t domain_size
     return 0;
 }
 
-EllRelaxContext *mg2di_ell_relax_alloc(size_t domain_size[2])
+EGSContext *mg2di_egs_alloc(enum EGSType type, size_t domain_size[2])
 {
-    EllRelaxContext *ctx;
-    EllRelaxInternal *priv;
+    EGSContext *ctx;
+    EGSInternal *priv;
     int ret;
 
     ctx = calloc(1, sizeof(*ctx));
@@ -576,11 +594,21 @@ EllRelaxContext *mg2di_ell_relax_alloc(size_t domain_size[2])
         goto fail;
     priv = ctx->priv;
 
+    switch (type) {
+    case EGS_SOLVER_RELAXATION:
+        ctx->solver_data = calloc(1, sizeof(EGSRelaxContext));
+        if (!ctx->solver_data)
+            goto fail;
+        break;
+    default: goto fail;
+    }
+    ctx->solver_type = type;
+
     if (!domain_size[0] || !domain_size[1] ||
         domain_size[0] > SIZE_MAX / domain_size[1])
         goto fail;
 
-    ret = ell_relax_arrays_alloc(ctx, domain_size);
+    ret = arrays_alloc(ctx, domain_size);
     if (ret < 0)
         goto fail;
 
@@ -589,17 +617,19 @@ EllRelaxContext *mg2di_ell_relax_alloc(size_t domain_size[2])
 
     return ctx;
 fail:
-    mg2di_ell_relax_free(&ctx);
+    mg2di_egs_free(&ctx);
     return NULL;
 }
 
-void mg2di_ell_relax_free(EllRelaxContext **pctx)
+void mg2di_egs_free(EGSContext **pctx)
 {
-    EllRelaxContext *ctx = *pctx;
+    EGSContext *ctx = *pctx;
 
     if (!ctx)
         return;
 
+    free(ctx->solver_data);
+
     free(ctx->priv->u_base);
     free(ctx->priv->rhs_base);
     free(ctx->priv->residual_base);
diff --git a/ell_relax.h b/ell_grid_solve.h
index 6a9ff2c..70782c8 100644
--- a/ell_relax.h
+++ b/ell_grid_solve.h
@@ -1,5 +1,5 @@
 /*
- * Relaxation solver for a 2nd order 2D linear PDE.
+ * Solver for a 2nd order 2D linear PDE using finite differencing on a grid.
  * Copyright 2018 Anton Khirnov <anton@khirnov.net>
  *
  * This program is free software: you can redistribute it and/or modify
@@ -16,8 +16,8 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-#ifndef MG2D_ELL_RELAX_H
-#define MG2D_ELL_RELAX_H
+#ifndef MG2D_ELL_GRID_SOLVE_H
+#define MG2D_ELL_GRID_SOLVE_H
 
 /**
  * The problem being solved is a linear partial differential
@@ -44,22 +44,50 @@
 #include "mg2d_boundary.h"
 #include "mg2d_constants.h"
 
-typedef struct EllRelaxInternal EllRelaxInternal;
+enum EGSType {
+    /**
+     * Solve the equation using relaxation.
+     *
+     * solver_data is EGSRelaxContext
+     * mg2di_egs_solve() does a single relaxation step
+     */
+    EGS_SOLVER_RELAXATION,
+};
+
+typedef struct EGSInternal EGSInternal;
+
+typedef struct EGSRelaxContext {
+    /**
+     * The time stepping factor in relaxation.
+     */
+    double relax_factor;
+
+    int64_t count_correct;
+    int64_t time_correct;
+} EGSRelaxContext;
 
-typedef struct EllRelaxContext {
+typedef struct EGSContext {
+    enum EGSType solver_type;
+
+    /**
+     * Solver type-specific data.
+     *
+     * Should be cast into the struct specified in documentation for this type.
+     */
+    void *solver_data;
     /**
      * Solver private data, not to be accessed in any way by the caller.
      */
-    EllRelaxInternal *priv;
+    EGSInternal *priv;
 
     /**
-     * The logging context, to be filled by the caller before mg2di_ell_relax_init().
+     * The logging context, to be filled by the caller before mg2di_egs_init().
      */
     MG2DLogger logger;
 
     /**
      * The thread pool used for execution. May be set by the caller before
-     * mg2di_ell_relax_init().
+     * mg2di_egs_init().
      */
     TPContext *tp;
 
@@ -69,41 +97,36 @@ typedef struct EllRelaxContext {
     int cpuflags;
 
     /**
-     * Size of the solver grid, set by mg2di_ell_relax_alloc().
+     * Size of the solver grid, set by mg2di_egs_alloc().
      * Read-only for the caller.
      */
     size_t domain_size[2];
 
     /**
      * Distance between the neighbouring grid points.
-     * Must be set by the caller before mg2di_ell_relax_init().
+     * Must be set by the caller before mg2di_egs_init().
      */
     double step[2];
 
     /**
      * Order of the finite difference operators used for approximating derivatives.
-     * Must be set by the caller before mg2di_ell_relax_init().
+     * Must be set by the caller before mg2di_egs_init().
      */
     size_t fd_stencil;
 
     /**
      * Boundary specification, indexed by MG2DBoundaryLoc.
-     * To be filled by the caller before mg2di_ell_relax_init().
+     * To be filled by the caller before mg2di_egs_init().
      */
     MG2DBoundary *boundaries[4];
 
     /**
-     * The time stepping factor in relaxation.
-     */
-    double relax_factor;
-
-    /**
      * Values of the unknown function.
      *
-     * Allocated by the solver in mg2di_ell_relax_alloc(), owned by the solver.
-     * Must be filled by the caller before mg2di_ell_relax_init() to set the
+     * Allocated by the solver in mg2di_egs_alloc(), owned by the solver.
+     * Must be filled by the caller before mg2di_egs_init() to set the
      * initial guess.
-     * Afterwards updated in mg2di_ell_relax_step().
+     * Afterwards updated in mg2di_egs_step().
      */
     double   *u;
     /**
@@ -114,8 +137,8 @@ typedef struct EllRelaxContext {
     /**
      * Values of the right-hand side.
      *
-     * Allocated by the solver in mg2di_ell_relax_alloc(), owned by the solver.
-     * Must be filled by the caller before mg2di_ell_relax_init().
+     * Allocated by the solver in mg2di_egs_alloc(), owned by the solver.
+     * Must be filled by the caller before mg2di_egs_init().
      */
     double    *rhs;
     /**
@@ -126,9 +149,9 @@ typedef struct EllRelaxContext {
     /**
      * Values of the right-hand side.
      *
-     * Allocated by the solver in mg2di_ell_relax_alloc(), owned by the solver.
-     * Read-only for the caller. Initialized after mg2di_ell_relax_init(),
-     * afterwards updated in mg2di_ell_relax_step().
+     * Allocated by the solver in mg2di_egs_alloc(), owned by the solver.
+     * Read-only for the caller. Initialized after mg2di_egs_init(),
+     * afterwards updated in mg2di_egs_step().
      */
     double   *residual;
     /**
@@ -144,8 +167,8 @@ typedef struct EllRelaxContext {
     /**
      * Coefficients C_{*} that define the differential equation.
      *
-     * Allocated by the solver in mg2di_ell_relax_alloc(), owned by the solver.
-     * Must be filled by the caller before mg2di_ell_relax_init().
+     * Allocated by the solver in mg2di_egs_alloc(), owned by the solver.
+     * Must be filled by the caller before mg2di_egs_init().
      */
     double   *diff_coeffs[MG2D_DIFF_COEFF_NB];
     /**
@@ -154,14 +177,11 @@ typedef struct EllRelaxContext {
     ptrdiff_t diff_coeffs_stride;
 
     /* timings */
-    int64_t count_correct;
-    int64_t time_correct;
-
     int64_t time_boundaries;
     int64_t count_boundaries;
     int64_t time_res_calc;
     int64_t count_res;
-} EllRelaxContext;
+} EGSContext;
 
 /**
  * Allocate the solver for the given domain size.
@@ -170,7 +190,7 @@ typedef struct EllRelaxContext {
  *
  * @return The solver context on success, NULL on failure.
  */
-EllRelaxContext *mg2di_ell_relax_alloc(size_t domain_size[2]);
+EGSContext *mg2di_egs_alloc(enum EGSType solver_type, size_t domain_size[2]);
 /**
  * Initialize the solver for use, after all the required fields are filled by
  * the caller.
@@ -181,17 +201,18 @@ EllRelaxContext *mg2di_ell_relax_alloc(size_t domain_size[2]);
  *
  * @return 0 on success, a negative error code on failure.
  */
-int              mg2di_ell_relax_init(EllRelaxContext *ctx);
+int              mg2di_egs_init(EGSContext *ctx);
 /**
  * Free the solver and write NULL to the provided pointer.
  */
-void             mg2di_ell_relax_free(EllRelaxContext **ctx);
+void             mg2di_egs_free(EGSContext **ctx);
 
 /**
- * Perform a single relaxation step.
+ * Solve the equation defined by the data filled in the context. Precise
+ * semantics depends on the solver type.
  *
  * @return 0 on success, a negative error code on failure.
  */
-int mg2di_ell_relax_step(EllRelaxContext *ctx);
+int mg2di_egs_solve(EGSContext *ctx);
 
-#endif /* MG2D_ELL_RELAX_H */
+#endif /* MG2D_ELL_GRID_SOLVE_H */
diff --git a/meson.build b/meson.build
index edd3b6f..fb20156 100644
--- a/meson.build
+++ b/meson.build
@@ -6,7 +6,7 @@ add_project_arguments('-D_XOPEN_SOURCE=700', language : 'c')
 lib_src = [
     'boundary.c',
     'cpu.c',
-    'ell_relax.c',
+    'ell_grid_solve.c',
     'log.c',
     'mg2d.c',
 ]
diff --git a/mg2d.c b/mg2d.c
index 3406030..9216af4 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -27,7 +27,7 @@
 
 #include "common.h"
 #include "cpu.h"
-#include "ell_relax.h"
+#include "ell_grid_solve.h"
 #include "log.h"
 #include "mg2d.h"
 #include "mg2d_boundary.h"
@@ -39,7 +39,7 @@
 typedef struct MG2DLevel {
     unsigned int  depth;
 
-    EllRelaxContext *solver;
+    EGSContext *solver;
 
     double   *prolong_tmp;
     ptrdiff_t prolong_tmp_stride;
@@ -86,8 +86,8 @@ static void log_relax_step(MG2DContext *ctx, MG2DLevel *level, const char *step_
               prefix, level->depth, step_desc, res_old, res_new, res_old / res_new);
 }
 
-static void mg_restrict_inject(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_stride,
-                               EllRelaxContext *fine, const double *src, ptrdiff_t src_stride)
+static void mg_restrict_inject(EGSContext *coarse, double *dst, ptrdiff_t dst_stride,
+                               EGSContext *fine, const double *src, ptrdiff_t src_stride)
 {
     for (size_t j = 0; j < coarse->domain_size[1]; j++)
         for (size_t i = 0; i < coarse->domain_size[0]; i++) {
@@ -101,8 +101,8 @@ static void mg_restrict_inject(EllRelaxContext *coarse, double *dst, ptrdiff_t d
         }
 }
 
-static void mg_restrict_fw(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_stride,
-                           EllRelaxContext *fine, const double *src, ptrdiff_t src_stride)
+static void mg_restrict_fw(EGSContext *coarse, double *dst, ptrdiff_t dst_stride,
+                           EGSContext *fine, const double *src, ptrdiff_t src_stride)
 {
     for (size_t j = 0; j < coarse->domain_size[1]; j++)
         for (size_t i = 0; i < coarse->domain_size[0]; i++) {
@@ -117,8 +117,8 @@ static void mg_restrict_fw(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_s
         }
 }
 
-static void mg_prolongate(EllRelaxContext *fine,         double *dst, ptrdiff_t dst_stride,
-                          EllRelaxContext *coarse, const double *src, ptrdiff_t src_stride)
+static void mg_prolongate(EGSContext *fine,         double *dst, ptrdiff_t dst_stride,
+                          EGSContext *coarse, const double *src, ptrdiff_t src_stride)
 {
     // for simplicity, this code writes one point over the domain size;
     // this allows us to avoid treating the boundary layers specially
@@ -196,8 +196,8 @@ static int mg_interp_bilinear_kernel(void *arg, unsigned int job_idx, unsigned i
 }
 
 static void mg_interp_bilinear(TPContext *tp,
-                               EllRelaxContext *ctx_dst, double *dst, ptrdiff_t dst_stride,
-                               EllRelaxContext *ctx_src, const double *src, ptrdiff_t src_stride)
+                               EGSContext *ctx_dst, double *dst, ptrdiff_t dst_stride,
+                               EGSContext *ctx_src, const double *src, ptrdiff_t src_stride)
 {
     InterpParamsBilin ip;
     ip.dst_domain_size[0] = ctx_dst->domain_size[0];
@@ -228,8 +228,8 @@ static void coarse_correct_task(void *arg, unsigned int job_idx, unsigned int th
 
 static void restrict_residual(MG2DContext *ctx, MG2DLevel *dst, MG2DLevel *src)
 {
-    EllRelaxContext *s_src = src->solver;
-    EllRelaxContext *s_dst = dst->solver;
+    EGSContext *s_src = src->solver;
+    EGSContext *s_dst = dst->solver;
     if (s_src->domain_size[0] == 2 * (s_dst->domain_size[0] - 1) + 1) {
         mg_restrict_fw(s_dst, s_dst->rhs, s_dst->rhs_stride,
                        s_src, s_src->residual, s_src->residual_stride);
@@ -242,8 +242,8 @@ static void restrict_residual(MG2DContext *ctx, MG2DLevel *dst, MG2DLevel *src)
 
 static void prolong_solution(MG2DContext *ctx, MG2DLevel *dst, MG2DLevel *src)
 {
-    EllRelaxContext *s_src = src->solver;
-    EllRelaxContext *s_dst = dst->solver;
+    EGSContext *s_src = src->solver;
+    EGSContext *s_dst = dst->solver;
     if (s_dst->domain_size[0] == 2 * (s_src->domain_size[0] - 1) + 1) {
         mg_prolongate(s_dst, dst->prolong_tmp, dst->prolong_tmp_stride,
                       s_src, s_src->u, s_src->u_stride);
@@ -264,7 +264,7 @@ static int mg_relax_step(MG2DContext *ctx, MG2DLevel *level, const char *step_de
 
     start = gettime();
 
-    ret = mg2di_ell_relax_step(level->solver);
+    ret = mg2di_egs_solve(level->solver);
 
     level->time_relax += gettime() - start;
 
@@ -288,7 +288,7 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level)
                 level->solver->domain_size[1]);
 
         /* re-init the current-level solver (re-calc the residual) */
-        ret = mg2di_ell_relax_init(level->solver);
+        ret = mg2di_egs_init(level->solver);
         if (ret < 0)
             return ret;
     }
@@ -344,7 +344,7 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level)
         /* re-init the current-level solver (re-calc the residual) */
         res_prev = level->solver->residual_max;
         start = gettime();
-        ret = mg2di_ell_relax_init(level->solver);
+        ret = mg2di_egs_init(level->solver);
         if (ret < 0)
             return ret;
         level->time_reinit += gettime() - start;
@@ -393,8 +393,8 @@ static void bnd_copy(MG2DBoundary *bdst, const double *src, ptrdiff_t src_stride
 
 static void restrict_diff_coeffs(MG2DContext *ctx, MG2DLevel *dst, MG2DLevel *src)
 {
-    EllRelaxContext *s_src = src->solver;
-    EllRelaxContext *s_dst = dst->solver;
+    EGSContext *s_src = src->solver;
+    EGSContext *s_dst = dst->solver;
     if (s_src->domain_size[0] == 2 * (s_dst->domain_size[0] - 1) + 1) {
         for (int i = 0; i < ARRAY_ELEMS(s_src->diff_coeffs); i++) {
             mg_restrict_inject(s_dst, s_dst->diff_coeffs[i], s_dst->diff_coeffs_stride,
@@ -453,7 +453,10 @@ static int mg_levels_init(MG2DContext *ctx)
         cur->solver->cpuflags = priv->cpuflags;
 
         cur->solver->fd_stencil = ctx->fd_stencil;
-        cur->solver->relax_factor = ctx->cfl_factor;
+        if (cur->solver->solver_type == EGS_SOLVER_RELAXATION) {
+            EGSRelaxContext *r = cur->solver->solver_data;
+            r->relax_factor = ctx->cfl_factor;
+        }
 
         prev = cur;
         cur  = cur->child;
@@ -484,7 +487,7 @@ int mg2d_solve(MG2DContext *ctx)
 {
     MG2DInternal *priv = ctx->priv;
     MG2DLevel    *root = priv->root;
-    EllRelaxContext *s_root = root->solver;
+    EGSContext *s_root = root->solver;
     int64_t time_start;
     double res_prev, res_cur;
     int ret;
@@ -501,7 +504,7 @@ int mg2d_solve(MG2DContext *ctx)
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ell_relax_init(s_root);
+    ret = mg2di_egs_init(s_root);
     if (ret < 0)
         return ret;
 
@@ -544,7 +547,7 @@ static void mg_level_free(MG2DLevel **plevel)
         return;
 
     free(level->prolong_tmp);
-    mg2di_ell_relax_free(&level->solver);
+    mg2di_egs_free(&level->solver);
 
     free(level);
     *plevel = NULL;
@@ -564,7 +567,8 @@ static MG2DLevel *mg_level_alloc(const size_t domain_size)
         goto fail;
     level->prolong_tmp_stride = domain_size + 1;
 
-    level->solver = mg2di_ell_relax_alloc((size_t [2]){domain_size, domain_size});
+    level->solver = mg2di_egs_alloc(EGS_SOLVER_RELAXATION,
+                                    (size_t [2]){domain_size, domain_size});
     if (!level->solver)
         goto fail;
 
@@ -723,9 +727,17 @@ void mg2d_print_stats(MG2DContext *ctx, const char *prefix)
               prefix, priv->count_solve, priv->time_solve / 1e6, priv->time_solve / 1e3 / priv->count_solve);
 
     while (level) {
+        EGSRelaxContext *r = NULL;
         int64_t level_total = level->time_relax + level->time_prolong + level->time_restrict +
                               level->time_correct + level->time_reinit;
-        int64_t relax_total = level->solver->time_correct + level->solver->time_res_calc + level->solver->time_boundaries;
+        int64_t relax_total = (r ? r->time_correct : 0) + level->solver->time_res_calc + level->solver->time_boundaries;
+
+        if (level->solver->solver_type == EGS_SOLVER_RELAXATION)
+            r = level->solver->solver_data;
+
+        level_total = level->time_relax + level->time_prolong + level->time_restrict +
+                              level->time_correct + level->time_reinit;
+        relax_total = (r ? r->time_correct : 0) + level->solver->time_res_calc + level->solver->time_boundaries;
 
         levels_total += level_total;
 
@@ -742,7 +754,7 @@ void mg2d_print_stats(MG2DContext *ctx, const char *prefix)
                   level->time_correct  * 100.0 / level_total,
                   level->time_reinit   * 100.0 / level_total,
                   level->solver->time_res_calc   * 100.0 / relax_total,
-                  level->solver->time_correct    * 100.0 / relax_total,
+                  (r ? r->time_correct : 0)    * 100.0 / relax_total,
                   level->solver->time_boundaries * 100.0 / relax_total
                   );
 
diff --git a/relax_test.c b/relax_test.c
index e885d55..d211e35 100644
--- a/relax_test.c
+++ b/relax_test.c
@@ -4,7 +4,7 @@
 #include <stdint.h>
 #include <string.h>
 
-#include "ell_relax.h"
+#include "ell_grid_solve.h"
 #include "log.h"
 #include "mg2d_boundary.h"
 #include "mg2d_constants.h"
@@ -49,7 +49,7 @@ static double findmax(double *arr, size_t len)
 
 int main(int argc, char **argv)
 {
-    EllRelaxContext *ctx;
+    EGSContext *ctx;
     long int log2N, log2maxiter;
     long int N, maxiter;
     double res_old, res_new;
@@ -71,7 +71,7 @@ int main(int argc, char **argv)
     N       = (1L << log2N) + 1;
     maxiter = 1L << log2maxiter;
 
-    ctx = mg2di_ell_relax_alloc((size_t [2]){N, N});
+    ctx = mg2di_egs_alloc(EGS_SOLVER_RELAXATION, (size_t [2]){N, N});
     if (!ctx) {
         fprintf(stderr, "Error allocating the solver context\n");
         return 1;
@@ -164,7 +164,7 @@ int main(int argc, char **argv)
                sizeof(*ctx->diff_coeffs[0]) * ctx->domain_size[0]);
     }
 
-    ret = mg2di_ell_relax_init(ctx);
+    ret = mg2di_egs_init(ctx);
     if (ret < 0) {
         fprintf(stderr, "Error initializing the solver context: %d\n", ret);
         ret = 1;
@@ -174,7 +174,7 @@ int main(int argc, char **argv)
     res_old = findmax(ctx->residual, ctx->residual_stride * ctx->domain_size[1]);
 
     for (int i = 0; i < maxiter; i++) {
-        ret = mg2di_ell_relax_step(ctx);
+        ret = mg2di_egs_solve(ctx);
         if (ret < 0) {
             fprintf(stderr, "Error during relaxation\n");
             ret = 1;
@@ -206,6 +206,6 @@ int main(int argc, char **argv)
     }
 
 fail:
-    mg2di_ell_relax_free(&ctx);
+    mg2di_egs_free(&ctx);
     return ret;
 }