ell_relax: rewrite residual calculation

The order-specific function now calculates the complete residual for an
entire row, instead of the derivatives of u for a single point.
This is done in preparation for upcoming SIMD.

1 files changed, 122 insertions, 45 deletions

diff --git a/ell_relax.c b/ell_relax.c
index aa791d8..8d4cb34 100644
--- a/ell_relax.c
+++ b/ell_relax.c
@@ -39,7 +39,11 @@ struct EllRelaxInternal {
 
     double *boundary_val_base[4];
 
-    void (*derivatives_calc)(EllRelaxContext *ctx, double *dst, ptrdiff_t idx, ptrdiff_t stride);
+    void (*residual_calc_line)(size_t linesize, double *dst,
+                               ptrdiff_t stride, const double *u, const double *rhs,
+                               const double * const diff_coeffs[ELL_RELAX_DIFF_COEFF_NB],
+                               const double *fd_factors);
+    double fd_factors[ELL_RELAX_DIFF_COEFF_NB];
     double relax_factor;
 };
 
@@ -65,59 +69,119 @@ static const struct {
     },
 };
 
+static const double fd_denoms[][ELL_RELAX_DIFF_COEFF_NB] = {
+    {
+        [ELL_RELAX_DIFF_COEFF_00] = 1.0,
+        [ELL_RELAX_DIFF_COEFF_10] = 2.0,
+        [ELL_RELAX_DIFF_COEFF_01] = 2.0,
+        [ELL_RELAX_DIFF_COEFF_20] = 1.0,
+        [ELL_RELAX_DIFF_COEFF_02] = 1.0,
+        [ELL_RELAX_DIFF_COEFF_11] = 4.0,
+    },
+    {
+        [ELL_RELAX_DIFF_COEFF_00] = 1.0,
+        [ELL_RELAX_DIFF_COEFF_10] = 12.0,
+        [ELL_RELAX_DIFF_COEFF_01] = 12.0,
+        [ELL_RELAX_DIFF_COEFF_20] = 12.0,
+        [ELL_RELAX_DIFF_COEFF_02] = 12.0,
+        [ELL_RELAX_DIFF_COEFF_11] = 144.0,
+    },
+};
+
 static void
-derivatives_calc_s1(EllRelaxContext *ctx, double *dst, ptrdiff_t idx, ptrdiff_t stride)
+derivatives_calc_s1(double *dst, const double *u, const double *fd_factors, ptrdiff_t stride)
 {
-    const double *u = ctx->u;
-    const double dx = ctx->step[0];
-    const double dy = ctx->step[1];
-
-    dst[ELL_RELAX_DIFF_COEFF_00] = u[idx];
-    dst[ELL_RELAX_DIFF_COEFF_10] = (u[idx + 1]      - u[idx - 1])      / (2.0 * dx);
-    dst[ELL_RELAX_DIFF_COEFF_01] = (u[idx + stride] - u[idx - stride]) / (2.0 * dy);
+    dst[ELL_RELAX_DIFF_COEFF_00] = u[0];
+    dst[ELL_RELAX_DIFF_COEFF_10] = (u[1]      - u[-1])      * fd_factors[ELL_RELAX_DIFF_COEFF_10];
+    dst[ELL_RELAX_DIFF_COEFF_01] = (u[stride] - u[-stride]) * fd_factors[ELL_RELAX_DIFF_COEFF_01];
 
-    dst[ELL_RELAX_DIFF_COEFF_20] = (u[idx + 1]      - 2.0 * u[idx] + u[idx - 1])      / (SQR(dx));
-    dst[ELL_RELAX_DIFF_COEFF_02] = (u[idx + stride] - 2.0 * u[idx] + u[idx - stride]) / (SQR(dy));
+    dst[ELL_RELAX_DIFF_COEFF_20] = (u[1]      - 2.0 * u[0] + u[-1])      * fd_factors[ELL_RELAX_DIFF_COEFF_20];
+    dst[ELL_RELAX_DIFF_COEFF_02] = (u[stride] - 2.0 * u[0] + u[-stride]) * fd_factors[ELL_RELAX_DIFF_COEFF_02];
 
-    dst[ELL_RELAX_DIFF_COEFF_11] = (u[idx + 1 + stride] - u[idx + stride - 1] - u[idx - stride + 1] + u[idx - stride - 1]) / (4.0 * dx * dy);
+    dst[ELL_RELAX_DIFF_COEFF_11] = (u[1 + stride] - u[stride - 1] - u[-stride + 1] + u[-stride - 1]) * fd_factors[ELL_RELAX_DIFF_COEFF_11];
 }
 
 static void
-derivatives_calc_s2(EllRelaxContext *ctx, double *dst, ptrdiff_t idx, ptrdiff_t stride)
+derivatives_calc_s2(double *dst, const double *u, const double *fd_factors, ptrdiff_t stride)
 {
-    const double *u = ctx->u;
-    const double dx = ctx->step[0];
-    const double dy = ctx->step[1];
-
-    dst[ELL_RELAX_DIFF_COEFF_00] = u[idx];
-    dst[ELL_RELAX_DIFF_COEFF_10] = (-1.0 * u[idx + 2] + 8.0 * u[idx + 1]      - 8.0 * u[idx - 1] + 1.0 * u[idx - 2])                             / (12.0 * dx);
-    dst[ELL_RELAX_DIFF_COEFF_01] = (-1.0 * u[idx + 2 * stride] + 8.0 * u[idx + stride]      - 8.0 * u[idx - stride] + 1.0 * u[idx - 2 * stride]) / (12.0 * dy);
-
-    dst[ELL_RELAX_DIFF_COEFF_20] = (-1.0 * u[idx + 2] + 16.0 * u[idx + 1] - 30.0 * u[idx] + 16.0 * u[idx - 1] - 1.0 * u[idx - 2])                             / (12.0 * SQR(dx));
-    dst[ELL_RELAX_DIFF_COEFF_02] = (-1.0 * u[idx + 2 * stride] + 16.0 * u[idx + stride] - 30.0 * u[idx] + 16.0 * u[idx - stride] - 1.0 * u[idx - 2 * stride]) / (12.0 * SQR(dy));
-
-    dst[ELL_RELAX_DIFF_COEFF_11] = ( 1.0 * u[idx + 2 + 2 * stride] - 8.0 * u[idx + 2 + stride] + 8.0 * u[idx + 2 - stride] - 1.0 * u[idx + 2 - 2 * stride]
-                                    -8.0 * u[idx + 1 + 2 * stride] + 64.0 * u[idx + 1 + stride] - 64.0 * u[idx + 1 - stride] + 8.0 * u[idx + 1 - 2 * stride]
-                                    +8.0 * u[idx - 1 + 2 * stride] - 64.0 * u[idx - 1 + stride] + 64.0 * u[idx - 1 - stride] - 8.0 * u[idx - 1 - 2 * stride]
-                                    -1.0 * u[idx - 2 + 2 * stride] + 8.0 * u[idx - 2 + stride] - 8.0 * u[idx - 2 - stride] + 1.0 * u[idx - 2 - 2 * stride])
-                                     / (144.0 * dx * dy);
+    const double val    = u[0];
+
+    const double valxp1 = u[ 1];
+    const double valxp2 = u[ 2];
+    const double valxm1 = u[-1];
+    const double valxm2 = u[-2];
+    const double valyp1 = u[ 1 * stride];
+    const double valyp2 = u[ 2 * stride];
+    const double valym1 = u[-1 * stride];
+    const double valym2 = u[-2 * stride];
+
+    const double valxp1yp1 = u[ 1 + 1 * stride];
+    const double valxp1yp2 = u[ 1 + 2 * stride];
+    const double valxp1ym1 = u[ 1 - 1 * stride];
+    const double valxp1ym2 = u[ 1 - 2 * stride];
+
+    const double valxp2yp1 = u[ 2 + 1 * stride];
+    const double valxp2yp2 = u[ 2 + 2 * stride];
+    const double valxp2ym1 = u[ 2 - 1 * stride];
+    const double valxp2ym2 = u[ 2 - 2 * stride];
+
+    const double valxm1yp1 = u[-1 + 1 * stride];
+    const double valxm1yp2 = u[-1 + 2 * stride];
+    const double valxm1ym1 = u[-1 - 1 * stride];
+    const double valxm1ym2 = u[-1 - 2 * stride];
+
+    const double valxm2yp1 = u[-2 + 1 * stride];
+    const double valxm2yp2 = u[-2 + 2 * stride];
+    const double valxm2ym1 = u[-2 - 1 * stride];
+    const double valxm2ym2 = u[-2 - 2 * stride];
+
+    dst[ELL_RELAX_DIFF_COEFF_00] = val;
+    dst[ELL_RELAX_DIFF_COEFF_10] = (-1.0 * valxp2 + 8.0 * valxp1      - 8.0 * valxm1 + 1.0 * valxm2) * fd_factors[ELL_RELAX_DIFF_COEFF_10];
+    dst[ELL_RELAX_DIFF_COEFF_01] = (-1.0 * valyp2 + 8.0 * valyp1      - 8.0 * valym1 + 1.0 * valym2) * fd_factors[ELL_RELAX_DIFF_COEFF_01];
+
+    dst[ELL_RELAX_DIFF_COEFF_20] = (-1.0 * valxp2 + 16.0 * valxp1 - 30.0 * val + 16.0 * valxm1 - 1.0 * valxm2) * fd_factors[ELL_RELAX_DIFF_COEFF_20];
+    dst[ELL_RELAX_DIFF_COEFF_02] = (-1.0 * valyp2 + 16.0 * valyp1 - 30.0 * val + 16.0 * valym1 - 1.0 * valym2) * fd_factors[ELL_RELAX_DIFF_COEFF_02];
+
+    dst[ELL_RELAX_DIFF_COEFF_11] = ( 1.0 * valxp2yp2 -  8.0 * valxp2yp1 +  8.0 * valxp2ym1 - 1.0 * valxp2ym2
+                                    -8.0 * valxp1yp2 + 64.0 * valxp1yp1 - 64.0 * valxp1ym1 + 8.0 * valxp1ym2
+                                    +8.0 * valxm1yp2 - 64.0 * valxm1yp1 + 64.0 * valxm1ym1 - 8.0 * valxm1ym2
+                                    -1.0 * valxm2yp2 +  8.0 * valxm2yp1 -  8.0 * valxm2ym1 + 1.0 * valxm2ym2) * fd_factors[ELL_RELAX_DIFF_COEFF_11];
 }
 
-static inline double residual_calc_kernel(EllRelaxContext *ctx, ptrdiff_t idx)
+static void residual_calc_line_s1_c(size_t linesize, double *dst,
+                                    ptrdiff_t stride, const double *u, const double *rhs,
+                                    const double * const diff_coeffs[ELL_RELAX_DIFF_COEFF_NB],
+                                    const double *fd_factors)
 {
-    EllRelaxInternal *priv = ctx->priv;
+    for (size_t i = 0; i < linesize; i++) {
+        double u_vals[ELL_RELAX_DIFF_COEFF_NB];
+        double res;
 
-    double u[ELL_RELAX_DIFF_COEFF_NB];
-    double res;
+        derivatives_calc_s1(u_vals, u + i, fd_factors, stride);
 
-    priv->derivatives_calc(ctx, u, idx, priv->stride);
+        res = -rhs[i];
+        for (int j = 0; j < ARRAY_ELEMS(u_vals); j++)
+            res += u_vals[j] * diff_coeffs[j][i];
+        dst[i] = res;
+    }
+}
 
-    res = -ctx->rhs[idx];
-    for (int i = 0; i < ARRAY_ELEMS(u); i++)
-        res += u[i] * ctx->diff_coeffs[i][idx];
+static void residual_calc_line_s2_c(size_t linesize, double *dst,
+                                    ptrdiff_t stride, const double *u, const double *rhs,
+                                    const double * const diff_coeffs[ELL_RELAX_DIFF_COEFF_NB],
+                                    const double *fd_factors)
+{
+    for (size_t i = 0; i < linesize; i++) {
+        double u_vals[ELL_RELAX_DIFF_COEFF_NB];
+        double res;
 
-    return res;
+        derivatives_calc_s2(u_vals, u + i, fd_factors, stride);
 
+        res = -rhs[i];
+        for (int j = 0; j < ARRAY_ELEMS(u_vals); j++)
+            res += u_vals[j] * diff_coeffs[j][i];
+        dst[i] = res;
+    }
 }
 
 static void residual_calc(EllRelaxContext *ctx)
@@ -126,11 +190,17 @@ static void residual_calc(EllRelaxContext *ctx)
     int64_t start;
 
     start = gettime();
-    for (int idx1 = 0; idx1 < ctx->domain_size[1]; idx1++)
-        for (int idx0 = 0; idx0 < ctx->domain_size[0]; idx0++) {
-            ptrdiff_t idx = idx1 * priv->stride + idx0;
-            ctx->residual[idx] = residual_calc_kernel(ctx, idx);
-        }
+    for (int idx1 = 0; idx1 < ctx->domain_size[1]; idx1++) {
+        const ptrdiff_t offset = idx1 * priv->stride;
+        const double *diff_coeffs[ELL_RELAX_DIFF_COEFF_NB];
+
+        for (int i = 0; i < ARRAY_ELEMS(diff_coeffs); i++)
+            diff_coeffs[i] = ctx->diff_coeffs[i] + offset;
+
+        priv->residual_calc_line(ctx->domain_size[0], ctx->residual + offset,
+                                 priv->stride, ctx->u + offset, ctx->rhs + offset,
+                                 diff_coeffs, priv->fd_factors);
+    }
 
     for (int i = 0; i < ARRAY_ELEMS(ctx->boundaries); i++) {
         const ptrdiff_t strides[2] = { 1, priv->stride };
@@ -274,10 +344,10 @@ int mg2di_ell_relax_init(EllRelaxContext *ctx)
 
     switch (ctx->fd_stencil) {
     case 1:
-        priv->derivatives_calc = derivatives_calc_s1;
+        priv->residual_calc_line = residual_calc_line_s1_c;
         break;
     case 2:
-        priv->derivatives_calc = derivatives_calc_s2;
+        priv->residual_calc_line = residual_calc_line_s2_c;
         break;
     default:
         mg2di_log(&ctx->logger, 0, "Invalid finite difference stencil: %zd\n",
@@ -295,6 +365,13 @@ int mg2di_ell_relax_init(EllRelaxContext *ctx)
     else
         priv->relax_factor = ctx->relax_factor;
 
+    priv->fd_factors[ELL_RELAX_DIFF_COEFF_00] = 1.0 /  fd_denoms[ctx->fd_stencil - 1][ELL_RELAX_DIFF_COEFF_00];
+    priv->fd_factors[ELL_RELAX_DIFF_COEFF_10] = 1.0 / (fd_denoms[ctx->fd_stencil - 1][ELL_RELAX_DIFF_COEFF_10] * ctx->step[0]);
+    priv->fd_factors[ELL_RELAX_DIFF_COEFF_01] = 1.0 / (fd_denoms[ctx->fd_stencil - 1][ELL_RELAX_DIFF_COEFF_01] * ctx->step[1]);
+    priv->fd_factors[ELL_RELAX_DIFF_COEFF_20] = 1.0 / (fd_denoms[ctx->fd_stencil - 1][ELL_RELAX_DIFF_COEFF_20] * SQR(ctx->step[0]));
+    priv->fd_factors[ELL_RELAX_DIFF_COEFF_02] = 1.0 / (fd_denoms[ctx->fd_stencil - 1][ELL_RELAX_DIFF_COEFF_02] * SQR(ctx->step[1]));
+    priv->fd_factors[ELL_RELAX_DIFF_COEFF_11] = 1.0 / (fd_denoms[ctx->fd_stencil - 1][ELL_RELAX_DIFF_COEFF_11] * ctx->step[0] * ctx->step[1]);
+
     for (int i = 0; i < ARRAY_ELEMS(ctx->boundaries); i++) {
         if (ctx->boundaries[i].type == ELL_RELAX_BC_TYPE_FIXDIFF) {
            for (int k = 0; k < ctx->domain_size[boundary_def[i].stride_idx]; k++)