ell_grid_solve: use ndarray API for allocating internal arrays

1 files changed, 37 insertions, 32 deletions

diff --git a/ell_grid_solve.c b/ell_grid_solve.c
index 550f305..117cbd1 100644
--- a/ell_grid_solve.c
+++ b/ell_grid_solve.c
@@ -34,6 +34,7 @@
 #include "mg2d_boundary.h"
 #include "mg2d_boundary_internal.h"
 #include "mg2d_constants.h"
+#include "ndarray.h"
 #include "residual_calc.h"
 
 static const double relax_factors[FD_STENCIL_MAX] = {
@@ -63,10 +64,10 @@ typedef struct EGSExactInternal {
 
 struct EGSInternal {
     ptrdiff_t stride;
-    double *u_base;
-    double *rhs_base;
-    double *residual_base;
-    double *diff_coeffs_base[MG2D_DIFF_COEFF_NB];
+    NDArray *u_base;
+    NDArray *rhs_base;
+    NDArray *residual_base;
+    NDArray *diff_coeffs_base[MG2D_DIFF_COEFF_NB];
 
     ptrdiff_t residual_calc_offset;
     size_t residual_calc_size[2];
@@ -670,38 +671,40 @@ static int arrays_alloc(EGSContext *ctx, const size_t domain_size[2])
 
     const size_t ghosts = FD_STENCIL_MAX;
     const size_t size_padded[2] = {
-        domain_size[0] + 2 * ghosts,
         domain_size[1] + 2 * ghosts,
+        domain_size[0] + 2 * ghosts,
     };
     const size_t stride = size_padded[0];
     const size_t start_offset = ghosts * stride + ghosts;
-    const size_t arr_size = size_padded[0] * size_padded[1];
+    const Slice slice[2] = { SLICE(ghosts, -ghosts, 1),
+                             SLICE(ghosts, -ghosts, 1) };
     int ret;
 
-    ret = posix_memalign((void**)&priv->u_base, 32, sizeof(*priv->u_base) * arr_size);
-    if (ret != 0)
-        return -ret;
-    ctx->u        = priv->u_base + start_offset;
-    ctx->u_stride = stride;
-
-    ret = posix_memalign((void**)&priv->rhs_base, 32, sizeof(*priv->rhs_base) * arr_size);
-    if (ret != 0)
-        return -ret;
-    ctx->rhs        = priv->rhs_base + start_offset;
-    ctx->rhs_stride = stride;
-
-    ret = posix_memalign((void**)&priv->residual_base, 32, sizeof(*priv->residual_base) * arr_size);
-    if (ret != 0)
-        return -ret;
-    memset(priv->residual_base, 0, sizeof(*priv->residual_base) * arr_size);
-    ctx->residual        = priv->residual_base + start_offset;
-    ctx->residual_stride = stride;
+    ret = mg2di_ndarray_alloc(&priv->u_base, 2, size_padded, 0);
+    if (ret < 0)
+        return ret;
+    ctx->u        = priv->u_base->data + start_offset;
+    ctx->u_stride = priv->u_base->stride[0];
+
+    ret = mg2di_ndarray_alloc(&priv->rhs_base, 2, size_padded, 0);
+    if (ret < 0)
+        return ret;
+    ctx->rhs        = priv->rhs_base->data + start_offset;
+    ctx->rhs_stride = priv->rhs_base->stride[0];
+
+    ret = mg2di_ndarray_alloc(&priv->residual_base, 2, size_padded,
+                              NDARRAY_ALLOC_ZERO);
+    if (ret < 0)
+        return ret;
+    ctx->residual        = priv->residual_base->data + start_offset;
+    ctx->residual_stride = priv->residual_base->stride[0];
 
     for (int i = 0; i < ARRAY_ELEMS(ctx->diff_coeffs); i++) {
-        priv->diff_coeffs_base[i] = calloc(arr_size, sizeof(*priv->diff_coeffs_base));
-        if (!priv->diff_coeffs_base[i])
-            return -ENOMEM;
-        ctx->diff_coeffs[i] = priv->diff_coeffs_base[i] + start_offset;
+        ret = mg2di_ndarray_alloc(&priv->diff_coeffs_base[i], 2,
+                                  size_padded, NDARRAY_ALLOC_ZERO);
+        if (ret < 0)
+            return ret;
+        ctx->diff_coeffs[i] = priv->diff_coeffs_base[i]->data + start_offset;
     }
     ctx->diff_coeffs_stride = stride;
 
@@ -817,11 +820,13 @@ void mg2di_egs_free(EGSContext **pctx)
         mg2di_bicgstab_context_free(&e->bicgstab);
     }
 
-    free(ctx->priv->u_base);
-    free(ctx->priv->rhs_base);
-    free(ctx->priv->residual_base);
+    mg2di_ndarray_free(&ctx->a_u);
+    mg2di_ndarray_free(&ctx->priv->u_base);
+
+    mg2di_ndarray_free(&ctx->priv->rhs_base);
+    mg2di_ndarray_free(&ctx->priv->residual_base);
     for (int i = 0; i < ARRAY_ELEMS(ctx->priv->diff_coeffs_base); i++)
-        free(ctx->priv->diff_coeffs_base[i]);
+        mg2di_ndarray_free(&ctx->priv->diff_coeffs_base[i]);
     for (int i = 0; i < ARRAY_ELEMS(ctx->boundaries); i++)
         mg2di_bc_free(&ctx->boundaries[i]);