mg2d: implement multicomponent solves for coarser levels

1 files changed, 50 insertions, 15 deletions

diff --git a/mg2d.c b/mg2d.c
index 1230dcf..0702a35 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -740,25 +740,60 @@ static int mg_level_partition(DomainGeometry **res, size_t domain_size,
 {
     DomainGeometry *dg;
 
-    dg = mg2di_dg_alloc(1);
-    if (!dg)
-        return -ENOMEM;
+    if (domain_size <= (1 << 7) + 1) {
+        dg = mg2di_dg_alloc(1);
+        if (!dg)
+            return -ENOMEM;
 
-    dg->domain_size[0] = domain_size;
-    dg->domain_size[1] = domain_size;
+        dg->domain_size[0] = domain_size;
+        dg->domain_size[1] = domain_size;
+
+        dg->components[0].interior.start[0] = 0;
+        dg->components[0].interior.start[1] = 0;
+        dg->components[0].interior.size[0]  = domain_size;
+        dg->components[0].interior.size[1]  = domain_size;
+
+        dg->components[0].exterior.start[0] = -FD_STENCIL_MAX;
+        dg->components[0].exterior.start[1] = -FD_STENCIL_MAX;
+        dg->components[0].exterior.size[0]  = domain_size + 2 * FD_STENCIL_MAX;
+        dg->components[0].exterior.size[1]  = domain_size + 2 * FD_STENCIL_MAX;
+
+        for (int i = 0; i < 4; i++)
+            dg->components[0].bnd_is_outer[i] = 1;
+    } else {
+        const double step_scale = (double)(domain_size - 1) / (src->domain_size[0] - 1);
+        dg = mg2di_dg_alloc(src->nb_components);
+        if (!dg)
+            return -ENOMEM;
 
-    dg->components[0].interior.start[0] = 0;
-    dg->components[0].interior.start[1] = 0;
-    dg->components[0].interior.size[0]  = domain_size;
-    dg->components[0].interior.size[1]  = domain_size;
+        dg->domain_size[0] = domain_size;
+        dg->domain_size[1] = domain_size;
 
-    dg->components[0].exterior.start[0] = -FD_STENCIL_MAX;
-    dg->components[0].exterior.start[1] = -FD_STENCIL_MAX;
-    dg->components[0].exterior.size[0]  = domain_size + 2 * FD_STENCIL_MAX;
-    dg->components[0].exterior.size[1]  = domain_size + 2 * FD_STENCIL_MAX;
+        // FIXME duplicated with restrict extents calculation
+        for (int comp = 0; comp < src->nb_components; comp++) {
+            const Rect *rect_src = &src->components[comp].interior;
+            DomainComponent *dc_dst = &dg->components[comp];
 
-    for (int i = 0; i < 4; i++)
-        dg->components[0].bnd_is_outer[i] = 1;
+            for (int dim = 0; dim < 2; dim++) {
+                ptrdiff_t start, end;
+
+                start = ceil(rect_src->start[dim] * step_scale);
+                end = ceil((rect_src->start[dim] + rect_src->size[dim]) * step_scale);
+                end = MIN(end, domain_size);
+
+                dc_dst->interior.start[dim] = start;
+                dc_dst->interior.size[dim]  = (end >= start) ? end - start : 0;
+
+                dc_dst->exterior.start[dim] = start ? start : -FD_STENCIL_MAX;
+                dc_dst->exterior.size[dim]  = dc_dst->interior.size[dim] + (end < domain_size ? 0 : FD_STENCIL_MAX) + (start ? 0 : FD_STENCIL_MAX);
+            }
+
+            dc_dst->bnd_is_outer[MG2D_BOUNDARY_0L] = dc_dst->interior.start[0] == 0;
+            dc_dst->bnd_is_outer[MG2D_BOUNDARY_1L] = dc_dst->interior.start[1] == 0;
+            dc_dst->bnd_is_outer[MG2D_BOUNDARY_0U] = (dc_dst->interior.start[0] + dc_dst->interior.size[0]) == domain_size;
+            dc_dst->bnd_is_outer[MG2D_BOUNDARY_1U] = (dc_dst->interior.start[1] + dc_dst->interior.size[1]) == domain_size;
+        }
+    }
 
     *res = dg;