mg2d: simplify and speed up prolongation

1 files changed, 14 insertions, 14 deletions

diff --git a/mg2d.c b/mg2d.c
index 8968306..8b70f40 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -119,26 +119,26 @@ static void mg_restrict_fw(EllRelaxContext *coarse, double *dst, ptrdiff_t dst_s
         }
 }
 
-static void mg_prolongate(EllRelaxContext *fine, double *dst, ptrdiff_t dst_stride,
+static void mg_prolongate(EllRelaxContext *fine,         double *dst, ptrdiff_t dst_stride,
                           EllRelaxContext *coarse, const double *src, ptrdiff_t src_stride)
 {
+    // for simplicity, this code writes one point over the domain size;
+    // this allows us to avoid treating the boundary layers specially
+    // dst is allocated with extra ghost zones for this purpose
     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_src = j * src_stride + i;
             const ptrdiff_t idx_dst = 2 * (j * dst_stride + i);
-            dst[idx_dst] = src[idx_src];
-        }
 
-    for (size_t j = 0; j < fine->domain_size[1]; j += 2)
-        for (size_t i = 1; i < fine->domain_size[0]; i += 2) {
-            const ptrdiff_t idx_dst = j * dst_stride + i;
-            dst[idx_dst] = 0.5 * (dst[idx_dst - 1] + dst[idx_dst + 1]);
-        }
+            const double src00 = src[idx_src];
+            const double src10 = src[idx_src + 1];
+            const double src01 = src[idx_src + src_stride];
+            const double src11 = src[idx_src + src_stride + 1];
 
-    for (size_t j = 1; j < fine->domain_size[1]; j += 2)
-        for (size_t i = 0; i < fine->domain_size[0]; i++) {
-            const ptrdiff_t idx = j * dst_stride + i;
-            dst[idx] = 0.5 * (dst[idx - dst_stride] + dst[idx + dst_stride]);
+            dst[idx_dst]                  = src00;
+            dst[idx_dst + 1]              = 0.5 * (src00 + src10);
+            dst[idx_dst + dst_stride]     = 0.5 * (src00 + src01);
+            dst[idx_dst + dst_stride + 1] = 0.25 * (src00 + src10 + src01 + src11);
         }
 }
 
@@ -501,10 +501,10 @@ static MG2DLevel *mg_level_alloc(const size_t domain_size)
     if (!level)
         return NULL;
 
-    ret = posix_memalign((void**)&level->prolong_tmp, 32, sizeof(*level->prolong_tmp) * SQR(domain_size));
+    ret = posix_memalign((void**)&level->prolong_tmp, 32, sizeof(*level->prolong_tmp) * SQR(domain_size + 1));
     if (ret != 0)
         goto fail;
-    level->prolong_tmp_stride = domain_size;
+    level->prolong_tmp_stride = domain_size + 1;
 
     level->solver = mg2di_ell_relax_alloc((size_t [2]){domain_size, domain_size});
     if (!level->solver)