tmptransfer_dim

3 files changed, 102 insertions, 26 deletions

diff --git a/ndarray.h b/ndarray.h
index db58254..ce08056 100644
--- a/ndarray.h
+++ b/ndarray.h
@@ -49,10 +49,43 @@ typedef struct NDArray {
 #define NDIDX2D(arr, x, y)    ((arr)->stride[0] * (y) + (arr)->stride[1] * (x))
 #define NDIDX3D(arr, x, y, z) ((arr)->stride[0] * (z) + (arr)->stride[1] * (y) + (arr)->stride[0] * (x))
 
+static inline ptrdiff_t mg2di_ndarray_vidx(NDArray *arr, va_list ap)
+{
+    ptrdiff_t ret = 0;
+    for (unsigned int i = 0; i < arr->dims; i++)
+        ret += arr->stride[i] * va_arg(ap, ptrdiff_t);
+    return ret;
+}
+
+static inline ptrdiff_t mg2di_ndarray_idx(NDArray *arr, ...)
+{
+    ptrdiff_t ret;
+    va_list ap;
+    va_start(ap, arr);
+    ret = mg2di_ndarray_vidx(arr, ap);
+    va_end(ap);
+    return ret;
+}
+
 #define NDPTR1D(arr, x)       ((arr)->data + NDIDX1D(arr, x))
 #define NDPTR2D(arr, x, y)    ((arr)->data + NDIDX2D(arr, x, y))
 #define NDPTR3D(arr, x, y, z) ((arr)->data + NDIDX2D(arr, x, y, z))
 
+static inline double *mg2di_ndarray_vdataptr(NDArray *arr, va_list ap)
+{
+    return arr->data + mg2di_ndarray_vidx(arr, ap);
+}
+
+static inline double *mg2di_ndarray_dataptr(NDArray *arr, ...)
+{
+    va_list ap;
+    double *ret;
+    va_start(ap, arr);
+    ret = mg2di_ndarray_vdataptr(arr, ap);
+    va_end(ap);
+    return ret;
+}
+
 int mg2di_ndarray_alloc(NDArray **result, unsigned int dims,
                         const size_t * const size, unsigned int flags);
 void mg2di_ndarray_free(NDArray **a);
diff --git a/transfer.c b/transfer.c
index 6db1ad2..c58b52d 100644
--- a/transfer.c
+++ b/transfer.c
@@ -52,8 +52,8 @@ typedef struct GridTransferLagrange {
     const NDArray *src;
     NDArray       *dst;
 
-    ptrdiff_t *idx[2];
-    double    *fact[2];
+    ptrdiff_t **idx;
+    double    **fact;
 } GridTransferLagrange;
 
 #if HAVE_EXTERNAL_ASM
@@ -153,6 +153,9 @@ static int transfer_lagrange_init(GridTransferContext *ctx)
     GridTransferLagrange *priv = ctx->priv;
     int ret;
 
+    if (ctx->nb_dims > 2)
+        return -ENOSYS;
+
     switch (ctx->op) {
     case GRID_TRANSFER_LAGRANGE_1:
         priv->stencil = 2;
@@ -189,7 +192,12 @@ static int transfer_lagrange_init(GridTransferContext *ctx)
         return -EINVAL;
     }
 
-    for (int dim = 0; dim < 2; dim++) {
+    priv->idx  = calloc(ctx->nb_dims, sizeof(*priv->idx));
+    priv->fact = calloc(ctx->nb_dims, sizeof(*priv->fact));
+    if (!priv->idx || !priv->fact)
+        return -ENOMEM;
+
+    for (int dim = 0; dim < ctx->nb_dims; dim++) {
         priv->idx[dim] = calloc(ctx->dst.size[dim], sizeof(*priv->idx[dim]));
         if (!priv->idx[dim])
             return -ENOMEM;
@@ -210,10 +218,14 @@ static void transfer_lagrange_free(GridTransferContext *ctx)
 {
     GridTransferLagrange *priv = ctx->priv;
 
-    for (int dim = 0; dim < 2; dim++) {
-        free(priv->idx[dim]);
-        free(priv->fact[dim]);
+    for (int dim = 0; dim < ctx->nb_dims; dim++) {
+        if (priv->idx)
+            free(priv->idx[dim]);
+        if (priv->fact)
+            free(priv->fact[dim]);
     }
+    free(priv->idx);
+    free(priv->fact);
 }
 
 static int transfer_interp_task(void *arg, unsigned int job_idx, unsigned int thread_idx)
@@ -325,7 +337,10 @@ static int transfer_fw_init(GridTransferContext *ctx)
 {
     GridTransferFW *fw = ctx->priv;
 
-    for (int i = 0; i < 2; i++) {
+    if (ctx->nb_dims != 2)
+        return -ENOSYS;
+
+    for (int i = 0; i < ctx->nb_dims; i++) {
         /* the destination must have twice the stepsize of the source */
         if (fabs(ctx->dst.step[i] / ctx->src.step[i] - 2.0) > 1e-14)
             return -EINVAL;
@@ -376,15 +391,18 @@ int mg2di_gt_transfer(GridTransferContext *ctx,
 {
     const GridTransfer *t = transfer_ops[ctx->op];
 
-    if (dst->dims != 2 || src->dims != 2 ||
-        dst->shape[0] != ctx->dst.size[0] || dst->shape[1] != ctx->dst.size[1] ||
-        src->shape[0] != ctx->src.size[0] || src->shape[1] != ctx->src.size[1])
+    if (dst->dims != ctx->nb_dims || src->dims != ctx->nb_dims)
         return -EINVAL;
 
+    for (int dir = 0; dir < ctx->nb_dims; dir++) {
+        if (dst->shape[dir] != ctx->dst.size[dir] || src->shape[dir] != ctx->src.size[dir])
+            return -EINVAL;
+    }
+
     return t->transfer(ctx, dst, src);
 }
 
-GridTransferContext *mg2di_gt_alloc(enum GridTransferOperator op)
+GridTransferContext *mg2di_gt_alloc(unsigned int nb_dims, enum GridTransferOperator op)
 {
     GridTransferContext *ctx;
     const GridTransfer *t;
@@ -397,16 +415,28 @@ GridTransferContext *mg2di_gt_alloc(enum GridTransferOperator op)
     if (!ctx)
         return NULL;
 
+    ctx->src.start = calloc(dims, sizeof(*ctx->src.start));
+    ctx->src.size  = calloc(dims, sizeof(*ctx->src.size));
+    ctx->src.step  = calloc(dims, sizeof(*ctx->src.step));
+    ctx->dst.start = calloc(dims, sizeof(*ctx->dst.start));
+    ctx->dst.size  = calloc(dims, sizeof(*ctx->dst.size));
+    ctx->dst.step  = calloc(dims, sizeof(*ctx->dst.step));
+    if (!ctx->src.start || !ctx->src.size || !ctx->src.step ||
+        !ctx->dst.start || !ctx->dst.size || !ctx->dst.step)
+        goto fail;
+    ctx->nb_dims = nb_dims;
+
     if (t->priv_data_size) {
         ctx->priv = calloc(1, t->priv_data_size);
-        if (!ctx->priv) {
-            free(ctx);
-            return NULL;
-        }
+        if (!ctx->priv)
+            goto fail;
     }
     ctx->op = op;
 
     return ctx;
+fail:
+    mg2di_gt_free(&ctx);
+    return NULL;
 }
 
 int mg2di_gt_init(GridTransferContext *ctx)
@@ -418,7 +448,7 @@ int mg2di_gt_init(GridTransferContext *ctx)
         return -EINVAL;
 
     /* check that the destination grid is inside the source one */
-    for (int i = 0; i < 2; i++) {
+    for (int i = 0; i < ctx->nb_dims; i++) {
         double src_start, src_end, dst_start, dst_end;
 
         if (ctx->src.step[i] <= 0.0 || ctx->dst.step[i] <= 0.0 ||
@@ -459,6 +489,13 @@ void mg2di_gt_free(GridTransferContext **pctx)
 
     free(ctx->priv);
 
+    free(ctx->src.start);
+    free(ctx->src.size);
+    free(ctx->src.step);
+    free(ctx->dst.start);
+    free(ctx->dst.size);
+    free(ctx->dst.step);
+
     free(ctx);
     *pctx = NULL;
 }
diff --git a/transfer.h b/transfer.h
index 0a10f3a..6e8ab8f 100644
--- a/transfer.h
+++ b/transfer.h
@@ -36,25 +36,26 @@ enum GridTransferOperator {
 };
 
 /**
- * A 2-dimensional grid with regular spacing.
+ * An N-dimensional grid with regular spacing.
  */
-typedef struct RegularGrid2D {
+typedef struct RegularGrid {
     /**
      * Indices of the grid origin, relative to coordinate origin.
-     * (o0, o1) = (start[0] * step[0], start[1] * step[1])
+     * For each dimension i=0..N-1:
+     *   o_i = start[i] * step[i]
      */
-    ptrdiff_t           start[2];
+    ptrdiff_t          *start;
 
     /**
      * Number of points in the domain.
      */
-    size_t              size[2];
+    size_t             *size;
 
     /**
      * Size of the step between neighbouring grid points.
      */
-    double              step[2];
-} RegularGrid2D;
+    double             *step;
+} RegularGrid;
 
 typedef struct GridTransferContext {
     /**
@@ -77,14 +78,19 @@ typedef struct GridTransferContext {
     enum GridTransferOperator op;
 
     /**
+     * Number of dimensions.
+     */
+    unsigned int nb_dims;
+
+    /**
      * Source grid geometry, must be filled by the caller.
      */
-    RegularGrid2D src;
+    RegularGrid src;
 
     /**
      * Destination grid geometry, must be filled by the caller.
      */
-    RegularGrid2D dst;
+    RegularGrid dst;
 } GridTransferContext;
 
 /**
@@ -92,7 +98,7 @@ typedef struct GridTransferContext {
  *
  * @return newly allocated transfer context on success, NULL on failure
  */
-GridTransferContext *mg2di_gt_alloc(enum GridTransferOperator op);
+GridTransferContext *mg2di_gt_alloc(unsigned int dims, enum GridTransferOperator op);
 
 /**
  * Initialize the tranfer context after all the parameters have been set.