mg2d: export the local component extents in public API

2 files changed, 29 insertions, 2 deletions

diff --git a/mg2d.c b/mg2d.c
index a0865c2..578dfe0 100644
--- a/mg2d.c
+++ b/mg2d.c
@@ -1430,6 +1430,12 @@ MG2DContext *mg2d_solver_alloc(size_t domain_size)
 
     ctx->domain_size = domain_size;
 
+    ctx->local_start[0] = 0;
+    ctx->local_start[1] = 0;
+
+    ctx->local_size[0] = domain_size;
+    ctx->local_size[1] = domain_size;
+
     return ctx;
 fail:
     mg2d_solver_free(&ctx);
@@ -1509,6 +1515,10 @@ MG2DContext *mg2d_solver_alloc_mpi(MPI_Comm comm, const size_t local_start[2],
     ctx->priv->dg              = dg;
     ctx->priv->local_component = rank;
     ctx->domain_size           = dg->domain_size[0];
+    ctx->local_start[0]        = local_start[0];
+    ctx->local_start[1]        = local_start[1];
+    ctx->local_size[0]         = local_size[0];
+    ctx->local_size[1]         = local_size[1];
 
     return ctx;
 fail:
@@ -1716,8 +1726,8 @@ int mg2d_init_guess(MG2DContext *ctx, const double *src,
         priv->transfer_init->src.step[0] = src_step[0];
         priv->transfer_init->src.step[1] = src_step[1];
 
-        priv->transfer_init->dst.size[0] = ctx->domain_size;
-        priv->transfer_init->dst.size[1] = ctx->domain_size;
+        priv->transfer_init->dst.size[0] = ctx->local_size[0];
+        priv->transfer_init->dst.size[1] = ctx->local_size[1];
         priv->transfer_init->dst.step[0] = ctx->step[0];
         priv->transfer_init->dst.step[1] = ctx->step[1];
 
diff --git a/mg2d.h b/mg2d.h
index 1db4ace..47d709e 100644
--- a/mg2d.h
+++ b/mg2d.h
@@ -162,6 +162,23 @@ typedef struct MG2DContext {
      * exact solve is performed.
      */
     size_t max_exact_size;
+
+    /**
+     * Indices of the lower left corner of the local domain in the full
+     * computational grid, for distributed solvers. Always { 0, 0 } for
+     * single-component solvers.
+     *
+     * Set by mg2d_solver_alloc[_mpi]().
+     */
+    ptrdiff_t local_start[2];
+
+    /**
+     * Number of points in the local component in each direction. Equal to
+     * { domain_size, domain_size } for single-component solvers.
+     *
+     * Set by mg2d_solver_alloc[_mpi]().
+     */
+    size_t    local_size[2];
 } MG2DContext;
 
 /**