mg2d: use injection for diff coeffs restriction

1 files changed, 21 insertions, 7 deletions

diff --git a/mg2d.c b/mg2d.c
index 7dacece..e019253 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -69,8 +69,23 @@ static void log_callback(MG2DLogger *log, int level, const char *fmt, va_list vl
     ctx->log_callback(ctx, level, fmt, vl);
 }
 
-static void mg_restrict(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_stride,
-                        EllRelaxContext *fine, const double *src, ptrdiff_t src_stride)
+static void mg_restrict_inject(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_stride,
+                               EllRelaxContext *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++) {
+            const ptrdiff_t idx_dst = j * dst_stride + i;
+
+            const size_t fine_i = i * 2;
+            const size_t fine_j = j * 2;
+            const ptrdiff_t idx_src = fine_j * src_stride + fine_i;
+
+            dst[idx_dst] = src[idx_src];
+        }
+}
+
+static void mg_restrict_fw(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_stride,
+                           EllRelaxContext *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++) {
@@ -80,7 +95,6 @@ static void mg_restrict(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_stri
             const size_t fine_j = j * 2;
             const ptrdiff_t idx_src = fine_j * src_stride + fine_i;
 
-            //dst[idx_dst] = src[idx_src];
             dst[idx_dst] = 0.25 * src[idx_src] + 0.125 * (src[idx_src + 1] + src[idx_src - 1] + src[idx_src + src_stride] + src[idx_src - src_stride])
                            + 0.0625 * (src[idx_src + 1 + src_stride] + src[idx_src - 1 + src_stride] + src[idx_src + 1 - src_stride] + src[idx_src - 1 - src_stride]);
         }
@@ -183,8 +197,8 @@ static int mg_solve_subgrid(MG2DContext *ctx, MG2DLevel *level)
             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);
+            mg_restrict_fw(level->child->solver, level->child->solver->rhs, level->child->solver->rhs_stride,
+                           level->solver, level->solver->residual, level->solver->residual_stride);
         }
 
 
@@ -242,8 +256,8 @@ static int mg_levels_init(MG2DContext *ctx)
                     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);
+                    mg_restrict_inject(cur->solver, cur->solver->diff_coeffs[i], cur->solver->diff_coeffs_stride,
+                                       prev->solver, prev->solver->diff_coeffs[i], prev->solver->diff_coeffs_stride);
                 }
             }
         }