egs: add support for MPI-based multi-component solves

1 files changed, 236 insertions, 0 deletions

diff --git a/relax_mpi_test.c b/relax_mpi_test.c
new file mode 100644
index 0000000..65888f6
--- /dev/null
+++ b/relax_mpi_test.c
@@ -0,0 +1,236 @@
+#include <math.h>
+#include <stdlib.h>
+#include <stdio.h>
+#include <stdint.h>
+#include <string.h>
+
+#include <mpi.h>
+
+#include "components.h"
+#include "ell_grid_solve.h"
+#include "log.h"
+#include "mg2d_boundary.h"
+#include "mg2d_constants.h"
+#include "ndarray.h"
+
+#define DOMAIN_SIZE 1.0
+#define FD_STENCIL 2
+
+#if 1
+static double sol(double x, double y)
+{
+    return sin(M_PI * x) * sin(M_PI * y);
+}
+static double sol_dx(double x, double y)
+{
+    return M_PI * cos(M_PI * x) * sin(M_PI * y);
+}
+static double sol_dy(double x, double y)
+{
+    return M_PI * sin(M_PI * x) * cos(M_PI * y);
+}
+static double sol_dxx(double x, double y)
+{
+    return -M_PI * M_PI * sol(x, y);
+}
+static double sol_dyy(double x, double y)
+{
+    return -M_PI * M_PI * sol(x, y);
+}
+static double sol_dxy(double x, double y)
+{
+    return M_PI * M_PI * cos(M_PI * x) * cos(M_PI * y);
+}
+#endif
+
+int main(int argc, char **argv)
+{
+    EGSContext *ctx;
+    DomainGeometry *dg = NULL;
+    DomainComponent *dc = NULL;
+    long int log2N, log2maxiter;
+    long int N, maxiter;
+    double res_old, res_new;
+    int ret = 0;
+
+    char processor_name[MPI_MAX_PROCESSOR_NAME];
+    int nb_processes, rank, processor_name_len;
+
+    if (argc < 3) {
+        fprintf(stderr, "Usage: %s <log2 N> <log2 maxiter>\n", argv[0]);
+        return 1;
+    }
+    log2N       = strtol(argv[1], NULL, 0);
+    log2maxiter = strtol(argv[2], NULL, 0);
+    if (log2N <= 0 || log2maxiter < 0 ||
+        log2N >= sizeof(N) * 8 || log2maxiter >= sizeof(maxiter) * 8) {
+        fprintf(stderr, "Invalid log2N/log2maxiter values: %ld/%ld\n",
+                log2N, log2maxiter);
+        return 1;
+    }
+
+    N       = (1L << log2N) + 1;
+    maxiter = 1L << log2maxiter;
+
+    MPI_Init(NULL, NULL);
+
+    MPI_Comm_size(MPI_COMM_WORLD, &nb_processes);
+    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+    MPI_Get_processor_name(processor_name, &processor_name_len);
+
+    fprintf(stderr, "This is process %d out of %d, running on %s\n",
+            rank, nb_processes, processor_name);
+
+    dg = mg2di_dg_alloc(nb_processes);
+    if (!dg) {
+        fprintf(stderr, "Error allocating domain geometry\n");
+        ret = 1;
+        goto fail;
+    }
+
+    dg->domain_size[0] = N;
+    dg->domain_size[1] = N;
+
+    for (unsigned int i = 0; i < dg->nb_components; i++) {
+        size_t patch_start, patch_end, patch_size_y;
+
+        patch_size_y = (N + nb_processes - 1) / nb_processes;
+        patch_start  = i * patch_size_y;
+        patch_end    = MIN((i + 1) * patch_size_y, N);
+        patch_size_y = patch_end - patch_start;
+        if (patch_size_y <= 0) {
+            fprintf(stderr, "Too many processes for grid size %ld: %d\n", N, nb_processes);
+            ret = 1;
+            goto fail;
+        }
+
+        dg->components[i].interior.start[0] = 0;
+        dg->components[i].interior.start[1] = patch_start;
+        dg->components[i].interior.size[0] = N;
+        dg->components[i].interior.size[1] = patch_size_y;
+
+        dg->components[i].exterior.start[0] = -FD_STENCIL;
+        dg->components[i].exterior.start[1] = i ? patch_start : -FD_STENCIL;
+        dg->components[i].exterior.size[0] = N + 2 * FD_STENCIL;
+        dg->components[i].exterior.size[1] = patch_size_y + ((i == 0) * FD_STENCIL) + ((i == nb_processes - 1) * FD_STENCIL);
+
+        dg->components[i].bnd_is_outer[MG2D_BOUNDARY_0L] = 1;
+        dg->components[i].bnd_is_outer[MG2D_BOUNDARY_0U] = 1;
+        dg->components[i].bnd_is_outer[MG2D_BOUNDARY_1L] = i == 0;
+        dg->components[i].bnd_is_outer[MG2D_BOUNDARY_1U] = i == nb_processes - 1;
+    }
+    dc = &dg->components[rank];
+
+    ctx = mg2di_egs_alloc_mpi(MPI_COMM_WORLD, dg);
+    if (!ctx) {
+        fprintf(stderr, "Error allocating the solver context\n");
+        return 1;
+    }
+
+    ctx->logger.log = mg2di_log_default_callback;
+
+    ctx->step[0] = DOMAIN_SIZE / (N - 1.0);
+    ctx->step[1] = DOMAIN_SIZE / (N - 1.0);
+
+    ctx->fd_stencil = FD_STENCIL;
+
+    for (int bnd_loc = 0; bnd_loc < ARRAY_ELEMS(ctx->boundaries); bnd_loc++) {
+        MG2DBoundary *bnd = ctx->boundaries[bnd_loc];
+        const int ci      = mg2d_bnd_coord_idx(bnd_loc);
+        const int bnd_dir = mg2d_bnd_out_dir(bnd_loc);
+
+        double coord[2];
+
+        if (!dc->bnd_is_outer[bnd_loc])
+            continue;
+
+        bnd->type = MG2D_BC_TYPE_FIXVAL;
+
+        memset(bnd->val, 0, ctx->domain_size[!ci] * sizeof(*bnd->val));
+
+        for (int j = 1; j < ctx->fd_stencil; j++) {
+            double *dst = bnd->val + j * bnd->val_stride;
+
+            coord[ci] = mg2d_bnd_is_upper(bnd_loc) * DOMAIN_SIZE + bnd_dir * j * ctx->step[ci];
+
+            for (ptrdiff_t k = -j; k < (ptrdiff_t)ctx->domain_size[!ci] + j; k++) {
+                coord[!ci] = (k + dc->interior.start[!ci]) * ctx->step[!ci];
+                dst[k] = sol(coord[0], coord[1]);
+            }
+        }
+    }
+
+    for (size_t y = 0; y < ctx->domain_size[1]; y++) {
+        const double y_coord = (y + dc->interior.start[1]) * ctx->step[1];
+
+        memset(NDPTR2D(ctx->u, 0, y), 0, sizeof(*ctx->u->data) * ctx->domain_size[0]);
+
+        for (size_t x = 1; x < ctx->domain_size[0]; x++) {
+            const double x_coord = x * ctx->step[0];
+
+            *NDPTR2D(ctx->diff_coeffs[MG2D_DIFF_COEFF_02], x, y) =  1.0;
+            *NDPTR2D(ctx->diff_coeffs[MG2D_DIFF_COEFF_20], x, y) =  1.0;
+            *NDPTR2D(ctx->diff_coeffs[MG2D_DIFF_COEFF_11], x, y) =  1.0;
+
+            *NDPTR2D(ctx->rhs, x, y) = sol_dxx(x_coord, y_coord) + sol_dyy(x_coord, y_coord) + sol_dxy(x_coord, y_coord);
+        }
+
+        memset(NDPTR2D(ctx->diff_coeffs[MG2D_DIFF_COEFF_00], 0, y), 0,
+               sizeof(*ctx->diff_coeffs[0]->data) * ctx->domain_size[0]);
+        memset(NDPTR2D(ctx->diff_coeffs[MG2D_DIFF_COEFF_01], 0, y), 0,
+               sizeof(*ctx->diff_coeffs[0]->data) * ctx->domain_size[0]);
+        memset(NDPTR2D(ctx->diff_coeffs[MG2D_DIFF_COEFF_10], 0, y), 0,
+               sizeof(*ctx->diff_coeffs[0]->data) * ctx->domain_size[0]);
+    }
+
+    ret = mg2di_egs_init(ctx, 0);
+    if (ret < 0) {
+        fprintf(stderr, "Error initializing the solver context: %d\n", ret);
+        ret = 1;
+        goto fail;
+    }
+
+    MPI_Reduce(&ctx->residual_max, &res_old, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD);
+
+    for (int i = 0; i < maxiter; i++) {
+        ret = mg2di_egs_solve(ctx, EGS_SOLVE_RELAXATION, 0);
+        if (ret < 0) {
+            fprintf(stderr, "Error during relaxation\n");
+            ret = 1;
+            goto fail;
+        }
+        MPI_Reduce(&ctx->residual_max, &res_new, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD);
+        if (rank == 0) {
+            if (res_new > 1e3) {
+                fprintf(stderr, "Diverged at step %d: %g -> %g\n", i, res_old, res_new);
+                goto fail;
+            }
+            fprintf(stderr, "Step %d: %g -> %g (%g)\n", i, res_old, res_new, res_old / res_new);
+            res_old = res_new;
+        }
+    }
+
+    {
+        double max_err = 0.0;
+
+        for (size_t y = 0; y < ctx->domain_size[1]; y++) {
+            const double y_coord = (y + dc->interior.start[1]) * ctx->step[1];
+
+            for (size_t x = 0; x < ctx->domain_size[0]; x++) {
+                const double x_coord = x * ctx->step[0];
+                double err = fabs(ctx->u->data[y * ctx->u->stride[0] + x] - sol(x_coord, y_coord));
+                if (err > max_err)
+                    max_err = err;
+            }
+        }
+        MPI_Reduce(rank ? &max_err : MPI_IN_PLACE, &max_err, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD);
+        if (rank == 0)
+            fprintf(stderr, "max(|solution - exact|): %g\n", max_err);
+    }
+
+fail:
+    mg2di_egs_free(&ctx);
+    mg2di_dg_free(&dg);
+    MPI_Finalize();
+    return ret;
+}