Add support for non-power-of-2 sized grids.

2 files changed, 83 insertions, 15 deletions

diff --git a/mg2d.c b/mg2d.c
index 1e605b8..73157f5 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -58,6 +58,11 @@ static double findmax(double *arr, size_t len)
     return ret;
 }
 
+static int is_power2(int n)
+{
+    return !(n & (n - 1));
+}
+
 static void log_callback(MG2DLogger *log, int level, const char *fmt, va_list vl)
 {
     MG2DContext *ctx = log->opaque;
@@ -104,6 +109,41 @@ static void mg_prolongate(EllRelaxContext *fine, double *dst, ptrdiff_t dst_stri
         }
 }
 
+static void mg_interp_bilinear(EllRelaxContext *ctx_dst, double *dst, ptrdiff_t dst_stride,
+                               EllRelaxContext *ctx_src, const double *src, ptrdiff_t src_stride)
+{
+    const double dx_dst = ctx_dst->step[0];
+    const double dy_dst = ctx_dst->step[1];
+    const double dx_src = ctx_src->step[0];
+    const double dy_src = ctx_src->step[1];
+    const double scalex = dx_dst / dx_src;
+    const double scaley = dy_dst / dy_src;
+
+    for (size_t idx1_dst = 0; idx1_dst < ctx_dst->domain_size[1]; idx1_dst++) {
+        const double y = idx1_dst * dy_dst;
+        const size_t idx1_src = idx1_dst * scaley;
+        const double y0 = idx1_src * dy_src;
+        const double y1 = y0 + dy_src;
+        const double fact_y0 = (y1 - y) / dy_src;
+        const double fact_y1 = (y - y0) / dy_src;
+
+        for (size_t idx0_dst = 0; idx0_dst < ctx_dst->domain_size[0]; idx0_dst++) {
+            const double x = idx0_dst * dx_dst;
+            const size_t idx0_src = idx0_dst * scalex;
+            const size_t idx_src = idx1_src * src_stride + idx0_src;
+            const double x0 = idx0_src * dx_src;
+            const double x1 = x0 + dx_src;
+            const double fact_x0 = (x1 - x) / dx_src;
+            const double fact_x1 = (x - x0) / dx_src;
+
+            const double val = fact_x0 * (fact_y0 * src[idx_src] + fact_y1 * src[idx_src + src_stride]) +
+                               fact_x1 * (fact_y0 * src[idx_src + 1] + fact_y1 * src[idx_src + 1 + src_stride]);
+
+            dst[idx1_dst * dst_stride + idx0_dst] = val;
+        }
+    }
+}
+
 static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level)
 {
     int ret;
@@ -139,8 +179,13 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level)
         }
 
         /* restrict the residual as to the coarser-level rhs */
-        mg_restrict(level->child->solver, level->child->solver->rhs, level->child->solver->rhs_stride,
-                    level->solver, level->solver->residual, level->solver->residual_stride);
+        if (!is_power2(level->solver->domain_size[0] - 1)) {
+            mg_interp_bilinear(level->child->solver, level->child->solver->rhs, level->child->solver->rhs_stride,
+                               level->solver, level->solver->residual, level->solver->residual_stride);
+        } else {
+            mg_restrict(level->child->solver, level->child->solver->rhs, level->child->solver->rhs_stride,
+                        level->solver, level->solver->residual, level->solver->residual_stride);
+        }
 
 
         /* solve on the coarser level */
@@ -149,8 +194,13 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level)
             return ret;
 
         /* prolongate the coarser-level correction */
-        mg_prolongate(level->solver, level->prolong_tmp, level->prolong_tmp_stride,
-                      level->child->solver, level->child->solver->u, level->child->solver->u_stride);
+        if (!is_power2(level->solver->domain_size[0] - 1)) {
+            mg_interp_bilinear(level->solver, level->prolong_tmp, level->prolong_tmp_stride,
+                               level->child->solver, level->child->solver->u, level->child->solver->u_stride);
+        } else {
+            mg_prolongate(level->solver, level->prolong_tmp, level->prolong_tmp_stride,
+                          level->child->solver, level->child->solver->u, level->child->solver->u_stride);
+        }
 
         /* apply the correction */
         for (size_t idx1 = 0; idx1 < level->solver->domain_size[1]; idx1++)
@@ -188,8 +238,13 @@ static int mg_levels_init(MG2DContext *ctx)
         /* Set the equation coefficients. */
         if (prev) {
             for (int i = 0; i < ARRAY_ELEMS(prev->solver->diff_coeffs); i++) {
-                mg_restrict(cur->solver, cur->solver->diff_coeffs[i], cur->solver->diff_coeffs_stride,
-                            prev->solver, prev->solver->diff_coeffs[i], prev->solver->diff_coeffs_stride);
+                if (!is_power2(prev->solver->domain_size[0] - 1)) {
+                    mg_interp_bilinear(cur->solver, cur->solver->diff_coeffs[i], cur->solver->diff_coeffs_stride,
+                                       prev->solver, prev->solver->diff_coeffs[i], prev->solver->diff_coeffs_stride);
+                } else {
+                    mg_restrict(cur->solver, cur->solver->diff_coeffs[i], cur->solver->diff_coeffs_stride,
+                                prev->solver, prev->solver->diff_coeffs[i], prev->solver->diff_coeffs_stride);
+                }
             }
         }
 
@@ -215,8 +270,13 @@ static int mg_levels_init(MG2DContext *ctx)
 
         cur->solver->logger = priv->logger;
 
-        cur->solver->step[0] = prev ? 2.0 * prev->solver->step[0] : ctx->step[0];
-        cur->solver->step[1] = prev ? 2.0 * prev->solver->step[1] : ctx->step[1];
+        if (!prev) {
+            cur->solver->step[0] = ctx->step[0];
+            cur->solver->step[1] = ctx->step[1];
+        } else {
+            cur->solver->step[0] = prev->solver->step[0] * ((double)(prev->solver->domain_size[0] - 1) / (cur->solver->domain_size[0] - 1));
+            cur->solver->step[1] = prev->solver->step[1] * ((double)(prev->solver->domain_size[1] - 1) / (cur->solver->domain_size[1] - 1));
+        }
 
         cur->solver->fd_stencil = ctx->fd_stencil;
 
@@ -303,6 +363,16 @@ fail:
     return NULL;
 }
 
+static int log2i(int n)
+{
+    int ret = 0;
+    while (n) {
+        ret++;
+        n >>= 1;
+    }
+    return ret - 1;
+}
+
 static int mg_levels_alloc(MG2DContext *ctx, size_t domain_size)
 {
     MG2DInternal *priv = ctx->priv;
@@ -313,8 +383,8 @@ static int mg_levels_alloc(MG2DContext *ctx, size_t domain_size)
         return -ENOMEM;
 
     cur = priv->root;
-    domain_size = (domain_size >> 1) + 1;
-    for (int i = 0; i < LEVELS_MAX && domain_size >= 4; i++) {
+    domain_size = (1 << log2i(domain_size - 2)) + 1;
+    for (int i = 0; i < LEVELS_MAX && domain_size > 4; i++) {
         cur->child = mg_level_alloc(domain_size);
         if (!cur->child)
             return -ENOMEM;
@@ -332,11 +402,10 @@ static void log_default_callback(const MG2DContext *ctx, int level, const char *
     vfprintf(stderr, fmt, vl);
 }
 
-MG2DContext *mg2d_solver_alloc(size_t log2_gridsize)
+MG2DContext *mg2d_solver_alloc(size_t domain_size)
 {
     MG2DContext *ctx;
     MG2DInternal *priv;
-    size_t domain_size;
     int ret;
 
     ctx = calloc(1, sizeof(*ctx));
@@ -363,9 +432,8 @@ MG2DContext *mg2d_solver_alloc(size_t log2_gridsize)
         ctx->boundaries[i]->val_len = ctx->domain_size;
     }
 
-    if (!((SIZE_MAX - 1) >> log2_gridsize))
+    if (SIZE_MAX / domain_size < domain_size)
         goto fail;
-    domain_size = (1 << log2_gridsize) + 1;
 
     ret = mg_levels_alloc(ctx, domain_size);
     if (ret < 0)
diff --git a/mg2d.h b/mg2d.h
index 3f9cbe7..405625a 100644
--- a/mg2d.h
+++ b/mg2d.h
@@ -226,7 +226,7 @@ typedef struct MG2DContext {
  *
  * @return The solver context on success, NULL on failure.
  */
-MG2DContext *mg2d_solver_alloc(size_t log2_gridsize);
+MG2DContext *mg2d_solver_alloc(size_t domain_size);
 /**
  * Solve the equation, after all the required fields have been filled by the
  * caller.