egs: make array allocation more readable

2 files changed, 60 insertions, 46 deletions

diff --git a/ell_grid_solve.c b/ell_grid_solve.c
index 61dc336..996d450 100644
--- a/ell_grid_solve.c
+++ b/ell_grid_solve.c
@@ -1067,88 +1067,102 @@ static int arrays_alloc(EGSContext *ctx, const size_t domain_size[2])
     EGSInternal *priv = ctx->priv;
     const DomainComponent *dc = &priv->dg->components[priv->local_component];
 
-    const size_t ghosts = FD_STENCIL_MAX;
-    const size_t size_padded[2] = {
-        domain_size[1] + 2 * ghosts,
-        domain_size[0] + 2 * ghosts,
-    };
-    const size_t size_dc[2] = {
-        size_padded[0] * MG2D_DIFF_COEFF_NB,
-        size_padded[1],
-    };
-    const Slice slice[2] = { SLICE(ghosts, -ghosts, 1),
-                             SLICE(ghosts, -ghosts, 1) };
-    Slice slice_exterior[2] = {
-        SLICE(dc->bnd_is_outer[MG2D_BOUNDARY_1L] ? 0 : ghosts,
-              dc->bnd_is_outer[MG2D_BOUNDARY_1U] ? 0 : -ghosts, 1),
-        SLICE(dc->bnd_is_outer[MG2D_BOUNDARY_0L] ? 0 : ghosts,
-              dc->bnd_is_outer[MG2D_BOUNDARY_0U] ? 0 : -ghosts, 1),
-    };
+    Slice slice_interior[2];
+    Slice slice_exterior[2];
+
+    size_t size_alloc[2];
+    size_t size_dc[2];
+    size_t extra_points[4];
+
     int ret;
 
-    ret = mg2di_ndarray_alloc(&priv->u_base, 2, size_padded, NDARRAY_ALLOC_ZERO);
+    for (int bnd_idx = 0; bnd_idx < 4; bnd_idx++)
+        extra_points[bnd_idx] = dc->bnd_is_outer[bnd_idx] ? 0 : FD_STENCIL_MAX;
+
+    size_alloc[0] = dc->exterior.size[1] + extra_points[MG2D_BOUNDARY_1L] + extra_points[MG2D_BOUNDARY_1U];
+    size_alloc[1] = dc->exterior.size[0] + extra_points[MG2D_BOUNDARY_0L] + extra_points[MG2D_BOUNDARY_0U];
+
+    slice_exterior[0].start = extra_points[MG2D_BOUNDARY_1L];
+    slice_exterior[0].end   = slice_exterior[0].start + dc->exterior.size[1];
+    slice_exterior[0].step  = 1;
+    slice_exterior[1].start = extra_points[MG2D_BOUNDARY_0L];
+    slice_exterior[1].end   = slice_exterior[1].start + dc->exterior.size[0];
+    slice_exterior[1].step  = 1;
+
+    slice_interior[0].start = slice_exterior[0].start + (dc->interior.start[1] - dc->exterior.start[1]);
+    slice_interior[0].end   = slice_interior[0].start + dc->interior.size[1];
+    slice_interior[0].step  = 1;
+    slice_interior[1].start = slice_exterior[1].start + (dc->interior.start[0] - dc->exterior.start[0]);
+    slice_interior[1].end   = slice_interior[1].start + dc->interior.size[0];
+    slice_interior[1].step  = 1;
+
+    ret = mg2di_ndarray_alloc(&priv->u_base,      2, size_alloc, NDARRAY_ALLOC_ZERO);
     if (ret < 0)
         return ret;
-
-    ret = mg2di_ndarray_slice(&ctx->u, priv->u_base, slice);
+    ret = mg2di_ndarray_alloc(&priv->u_next_base, 2, size_alloc, NDARRAY_ALLOC_ZERO);
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_slice(&ctx->u_exterior, priv->u_base, slice_exterior);
+    ret = mg2di_ndarray_slice(&ctx->u,       priv->u_base,      slice_interior);
     if (ret < 0)
         return ret;
-
-    ret = mg2di_ndarray_alloc(&priv->u_next_base, 2, size_padded, NDARRAY_ALLOC_ZERO);
+    ret = mg2di_ndarray_slice(&priv->u_next, priv->u_next_base, slice_interior);
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_slice(&priv->u_next, priv->u_next_base, slice);
+    ret = mg2di_ndarray_slice(&ctx->u_exterior,       priv->u_base,      slice_exterior);
     if (ret < 0)
         return ret;
-
     ret = mg2di_ndarray_slice(&priv->u_next_exterior, priv->u_next_base, slice_exterior);
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_alloc(&priv->rhs_base, 2, size_padded, 0);
+    ret = mg2di_ndarray_alloc(&priv->rhs_base, 2, size_alloc, 0);
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_slice(&ctx->rhs, priv->rhs_base, slice);
+    ret = mg2di_ndarray_slice(&ctx->rhs, priv->rhs_base, slice_interior);
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_alloc(&priv->residual_base, 2, size_padded,
+    ret = mg2di_ndarray_alloc(&priv->residual_base, 2, size_alloc,
                               NDARRAY_ALLOC_ZERO);
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_slice(&ctx->residual, priv->residual_base, slice);
+    ret = mg2di_ndarray_slice(&ctx->residual, priv->residual_base, slice_interior);
     if (ret < 0)
         return ret;
 
+    size_dc[0]    = size_alloc[0] * MG2D_DIFF_COEFF_NB;
+    size_dc[1]    = size_alloc[1];
+
     ret = mg2di_ndarray_alloc(&priv->diff_coeffs_public_base, 2,
                               size_dc, NDARRAY_ALLOC_ZERO);
     if (ret < 0)
         return ret;
-
-    for (int i = 0; i < ARRAY_ELEMS(ctx->diff_coeffs); i++) {
-        const Slice slice_dc[2] = { SLICE(i * size_padded[0] + ghosts, (i + 1) * size_padded[0] -ghosts, 1),
-                                    SLICE(ghosts, -ghosts, 1) };
-        ret = mg2di_ndarray_slice(&ctx->diff_coeffs[i], priv->diff_coeffs_public_base, slice_dc);
-        if (ret < 0)
-            return ret;
-    }
-
     ret = mg2di_ndarray_alloc(&priv->diff_coeffs_base, 2,
                               size_dc, NDARRAY_ALLOC_ZERO);
     if (ret < 0)
         return ret;
 
-    for (int i = 0; i < ARRAY_ELEMS(priv->diff_coeffs); i++) {
-        const Slice slice_dc[2] = { SLICE(i * size_padded[0] + ghosts, (i + 1) * size_padded[0] -ghosts, 1),
-                                    SLICE(ghosts, -ghosts, 1) };
-        ret = mg2di_ndarray_slice(&priv->diff_coeffs[i], priv->diff_coeffs_base, slice_dc);
+    for (int i = 0; i < ARRAY_ELEMS(ctx->diff_coeffs); i++) {
+        const Slice slice_dc[2] = { SLICE(i * size_alloc[0], (i + 1) * size_alloc[0], 1), SLICE_NULL };
+        NDArray *tmp;
+
+        ret = mg2di_ndarray_slice(&tmp, priv->diff_coeffs_public_base, slice_dc);
+        if (ret < 0)
+            return ret;
+        ret = mg2di_ndarray_slice(&ctx->diff_coeffs[i], tmp, slice_interior);
+        mg2di_ndarray_free(&tmp);
+        if (ret < 0)
+            return ret;
+
+        ret = mg2di_ndarray_slice(&tmp, priv->diff_coeffs_base, slice_dc);
+        if (ret < 0)
+            return ret;
+        ret = mg2di_ndarray_slice(&priv->diff_coeffs[i], tmp, slice_interior);
+        mg2di_ndarray_free(&tmp);
         if (ret < 0)
             return ret;
     }
diff --git a/relax_mpi_test.c b/relax_mpi_test.c
index 65888f6..75e76ef 100644
--- a/relax_mpi_test.c
+++ b/relax_mpi_test.c
@@ -109,10 +109,10 @@ int main(int argc, char **argv)
         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].exterior.start[0] = -FD_STENCIL_MAX;
+        dg->components[i].exterior.start[1] = i ? patch_start : -FD_STENCIL_MAX;
+        dg->components[i].exterior.size[0] = N + 2 * FD_STENCIL_MAX;
+        dg->components[i].exterior.size[1] = patch_size_y + ((i == 0) * FD_STENCIL_MAX) + ((i == nb_processes - 1) * FD_STENCIL_MAX);
 
         dg->components[i].bnd_is_outer[MG2D_BOUNDARY_0L] = 1;
         dg->components[i].bnd_is_outer[MG2D_BOUNDARY_0U] = 1;