Line data Source code
1 : /**
2 : * @file energy_solver_avx2.c
3 : * @brief AVX2-vectorized energy equation solver (advection-diffusion step)
4 : *
5 : * Solves: dT/dt + u*nabla(T) = alpha * nabla^2(T) + Q
6 : * using explicit Euler time integration and central differences.
7 : *
8 : * Same numerics as the scalar reference (energy/cpu/energy_solver.c). The
9 : * interior stencil is vectorized along the unit-stride i direction (4 doubles
10 : * per AVX2 register) with a scalar remainder tail; the j loop is parallelized
11 : * with OpenMP. When AVX2 is unavailable at compile time the whole interior is
12 : * processed with the scalar path, so this file always provides a correct
13 : * implementation. The optional heat source Q is applied in a unified scalar
14 : * pass (only when a callback is configured), keeping the vector path branchless.
15 : *
16 : * Branch-free 3D: when nz==1, stride_z=0 and inv_2dz/inv_dz2=0.0 cause all
17 : * z-terms to vanish, producing identical results to a 2D code path.
18 : */
19 :
20 : #define _USE_MATH_DEFINES
21 :
22 : #include "../energy_solver_internal.h"
23 :
24 : #include "cfd/core/indexing.h"
25 : #include "cfd/core/memory.h"
26 :
27 : #include <math.h>
28 : #include <stddef.h>
29 : #include <string.h>
30 :
31 : #ifdef _OPENMP
32 : #include <omp.h>
33 : #endif
34 :
35 : #if defined(CFD_HAS_AVX2)
36 : #include <immintrin.h>
37 : #define USE_AVX 1
38 : #else
39 : #define USE_AVX 0
40 : #endif
41 :
42 325 : cfd_status_t energy_step_explicit_avx2_with_workspace(
43 : flow_field* field, const grid* grid,
44 : const ns_solver_params_t* params,
45 : double dt, double time,
46 : double* T_workspace, size_t workspace_size) {
47 325 : if (!field || !grid || !params) {
48 0 : cfd_set_error(CFD_ERROR_INVALID,
49 : "energy_solver_avx2: field, grid, and params must be non-NULL");
50 0 : return CFD_ERROR_INVALID;
51 : }
52 325 : if (!field->T) {
53 0 : cfd_set_error(CFD_ERROR_INVALID,
54 : "energy_solver_avx2: missing temperature field");
55 0 : return CFD_ERROR_INVALID;
56 : }
57 :
58 : /* Skip when energy equation is disabled */
59 325 : if (params->alpha <= 0.0) {
60 : return CFD_SUCCESS;
61 : }
62 :
63 20 : size_t nx = field->nx;
64 20 : size_t ny = field->ny;
65 20 : size_t nz = field->nz;
66 :
67 20 : if (!grid->dx || !grid->dy || nx < 3 || ny < 3) {
68 0 : cfd_set_error(CFD_ERROR_INVALID,
69 : "energy_solver_avx2: grid too small or missing dx/dy");
70 0 : return CFD_ERROR_INVALID;
71 : }
72 20 : size_t plane = nx * ny;
73 20 : size_t total = plane * nz;
74 20 : double alpha = params->alpha;
75 :
76 : /* Validate uniform spacing (central-difference stencil assumes it) */
77 20 : const double dx0 = grid->dx[0];
78 20 : const double dy0 = grid->dy[0];
79 : {
80 20 : const double tol_x = 1e-12 * fmax(1.0, fabs(dx0));
81 160 : for (size_t i = 1; i < nx - 1; i++) {
82 140 : if (fabs(grid->dx[i] - dx0) > tol_x) {
83 0 : cfd_set_error(CFD_ERROR_UNSUPPORTED,
84 : "energy_solver_avx2: non-uniform dx not supported");
85 0 : return CFD_ERROR_UNSUPPORTED;
86 : }
87 : }
88 20 : const double tol_y = 1e-12 * fmax(1.0, fabs(dy0));
89 160 : for (size_t j = 1; j < ny - 1; j++) {
90 140 : if (fabs(grid->dy[j] - dy0) > tol_y) {
91 0 : cfd_set_error(CFD_ERROR_UNSUPPORTED,
92 : "energy_solver_avx2: non-uniform dy not supported");
93 0 : return CFD_ERROR_UNSUPPORTED;
94 : }
95 : }
96 : }
97 20 : if (nz > 1) {
98 0 : if (!grid->dz) {
99 0 : cfd_set_error(CFD_ERROR_INVALID,
100 : "energy_solver_avx2: missing dz for 3D energy solve");
101 0 : return CFD_ERROR_INVALID;
102 : }
103 0 : const double dz0 = grid->dz[0];
104 0 : const double tol_z = 1e-12 * fmax(1.0, fabs(dz0));
105 0 : for (size_t k = 1; k < nz - 1; k++) {
106 0 : if (fabs(grid->dz[k] - dz0) > tol_z) {
107 0 : cfd_set_error(CFD_ERROR_UNSUPPORTED,
108 : "energy_solver_avx2: non-uniform dz not supported");
109 0 : return CFD_ERROR_UNSUPPORTED;
110 : }
111 : }
112 : }
113 :
114 : /* Precomputed constants for uniform-grid stencil */
115 20 : double inv_2dx = 1.0 / (2.0 * dx0);
116 20 : double inv_2dy = 1.0 / (2.0 * dy0);
117 20 : double inv_dx2 = 1.0 / (dx0 * dx0);
118 20 : double inv_dy2 = 1.0 / (dy0 * dy0);
119 :
120 : /* Branch-free 3D constants */
121 20 : size_t stride_z = (nz > 1) ? plane : 0;
122 20 : size_t k_start = (nz > 1) ? 1 : 0;
123 20 : size_t k_end = (nz > 1) ? (nz - 1) : 1;
124 20 : double inv_2dz = (nz > 1 && grid->dz) ? 1.0 / (2.0 * grid->dz[0]) : 0.0;
125 0 : double inv_dz2 = (nz > 1 && grid->dz) ? 1.0 / (grid->dz[0] * grid->dz[0]) : 0.0;
126 :
127 : /* Use caller's workspace or allocate internally */
128 20 : int owns_buffer = 0;
129 20 : double* T_new;
130 20 : if (T_workspace && workspace_size >= total) {
131 : T_new = T_workspace;
132 : } else {
133 0 : T_new = (double*)cfd_calloc(total, sizeof(double));
134 0 : if (!T_new) {
135 : return CFD_ERROR_NOMEM;
136 : }
137 : owns_buffer = 1;
138 : }
139 20 : memcpy(T_new, field->T, total * sizeof(double));
140 :
141 20 : const double* T = field->T;
142 20 : const double* u = field->u;
143 20 : const double* v = field->v;
144 20 : const double* w = field->w;
145 :
146 : /* Number of interior columns processed in 4-wide vector groups (i=1..nx-2) */
147 20 : size_t n_interior = nx - 2;
148 20 : size_t i_tail = 1 + (n_interior / 4) * 4;
149 :
150 : #if USE_AVX
151 : const __m256d inv_2dx_v = _mm256_set1_pd(inv_2dx);
152 : const __m256d inv_2dy_v = _mm256_set1_pd(inv_2dy);
153 : const __m256d inv_2dz_v = _mm256_set1_pd(inv_2dz);
154 : const __m256d inv_dx2_v = _mm256_set1_pd(inv_dx2);
155 : const __m256d inv_dy2_v = _mm256_set1_pd(inv_dy2);
156 : const __m256d inv_dz2_v = _mm256_set1_pd(inv_dz2);
157 : const __m256d alpha_v = _mm256_set1_pd(alpha);
158 : const __m256d dt_v = _mm256_set1_pd(dt);
159 : const __m256d two_v = _mm256_set1_pd(2.0);
160 : #endif
161 :
162 40 : for (size_t k = k_start; k < k_end; k++) {
163 20 : size_t k_offset = k * stride_z;
164 20 : int j;
165 : #ifdef _OPENMP
166 20 : #pragma omp parallel for schedule(static)
167 : #endif
168 : for (j = 1; j < (int)ny - 1; j++) {
169 : size_t row = k_offset + (size_t)j * nx;
170 :
171 : #if USE_AVX
172 : for (size_t i = 1; i < i_tail; i += 4) {
173 : size_t idx = row + i;
174 :
175 : __m256d T_c = _mm256_loadu_pd(&T[idx]);
176 : __m256d uc = _mm256_loadu_pd(&u[idx]);
177 : __m256d vc = _mm256_loadu_pd(&v[idx]);
178 : __m256d wc = _mm256_loadu_pd(&w[idx]);
179 :
180 : __m256d T_xp = _mm256_loadu_pd(&T[idx + 1]);
181 : __m256d T_xm = _mm256_loadu_pd(&T[idx - 1]);
182 : __m256d T_yp = _mm256_loadu_pd(&T[idx + nx]);
183 : __m256d T_ym = _mm256_loadu_pd(&T[idx - nx]);
184 : __m256d T_zp = _mm256_loadu_pd(&T[idx + stride_z]);
185 : __m256d T_zm = _mm256_loadu_pd(&T[idx - stride_z]);
186 :
187 : /* Advection: u*dT/dx + v*dT/dy + w*dT/dz */
188 : __m256d dT_dx = _mm256_mul_pd(_mm256_sub_pd(T_xp, T_xm), inv_2dx_v);
189 : __m256d dT_dy = _mm256_mul_pd(_mm256_sub_pd(T_yp, T_ym), inv_2dy_v);
190 : __m256d dT_dz = _mm256_mul_pd(_mm256_sub_pd(T_zp, T_zm), inv_2dz_v);
191 : __m256d adv = _mm256_add_pd(
192 : _mm256_add_pd(_mm256_mul_pd(uc, dT_dx), _mm256_mul_pd(vc, dT_dy)),
193 : _mm256_mul_pd(wc, dT_dz));
194 :
195 : /* Diffusion: alpha * (d2T/dx2 + d2T/dy2 + d2T/dz2) */
196 : __m256d d2x = _mm256_mul_pd(
197 : _mm256_sub_pd(_mm256_add_pd(T_xp, T_xm), _mm256_mul_pd(two_v, T_c)), inv_dx2_v);
198 : __m256d d2y = _mm256_mul_pd(
199 : _mm256_sub_pd(_mm256_add_pd(T_yp, T_ym), _mm256_mul_pd(two_v, T_c)), inv_dy2_v);
200 : __m256d d2z = _mm256_mul_pd(
201 : _mm256_sub_pd(_mm256_add_pd(T_zp, T_zm), _mm256_mul_pd(two_v, T_c)), inv_dz2_v);
202 : __m256d diff = _mm256_mul_pd(alpha_v,
203 : _mm256_add_pd(_mm256_add_pd(d2x, d2y), d2z));
204 :
205 : /* Explicit Euler update (heat source added in the scalar pass) */
206 : __m256d dT = _mm256_mul_pd(dt_v, _mm256_sub_pd(diff, adv));
207 : _mm256_storeu_pd(&T_new[idx], _mm256_add_pd(T_c, dT));
208 : }
209 : #endif
210 : /* Scalar remainder (and the full interior when AVX2 is disabled) */
211 : for (size_t i = (USE_AVX ? i_tail : 1); i < nx - 1; i++) {
212 : size_t idx = row + i;
213 :
214 : double T_cs = T[idx];
215 : double dT_dx = (T[idx + 1] - T[idx - 1]) * inv_2dx;
216 : double dT_dy = (T[idx + nx] - T[idx - nx]) * inv_2dy;
217 : double dT_dz = (T[idx + stride_z] - T[idx - stride_z]) * inv_2dz;
218 : double adv = u[idx] * dT_dx + v[idx] * dT_dy + w[idx] * dT_dz;
219 :
220 : double d2x = (T[idx + 1] - 2.0 * T_cs + T[idx - 1]) * inv_dx2;
221 : double d2y = (T[idx + nx] - 2.0 * T_cs + T[idx - nx]) * inv_dy2;
222 : double d2z = (T[idx + stride_z] - 2.0 * T_cs + T[idx - stride_z]) * inv_dz2;
223 : double diff = alpha * (d2x + d2y + d2z);
224 :
225 : T_new[idx] = T_cs + dt * (diff - adv);
226 : }
227 :
228 : /* Heat source term (scalar; only when a callback is configured) */
229 : if (params->heat_source_func) {
230 : double y = grid->y[j];
231 : double z = (nz > 1 && grid->z) ? grid->z[k] : 0.0;
232 : for (size_t i = 1; i < nx - 1; i++) {
233 : size_t idx = row + i;
234 : double Q = params->heat_source_func(grid->x[i], y, z, time,
235 : params->heat_source_context);
236 : T_new[idx] += dt * Q;
237 : }
238 : }
239 : }
240 : }
241 :
242 : /* Check for NaN/Inf */
243 20 : int has_nan = 0;
244 20 : ptrdiff_t total_int = (ptrdiff_t)total;
245 20 : ptrdiff_t n;
246 : #ifdef _OPENMP
247 20 : #pragma omp parallel for reduction(| : has_nan) schedule(static)
248 : #endif
249 : for (n = 0; n < total_int; n++) {
250 : if (!isfinite(T_new[n])) {
251 : has_nan = 1;
252 : }
253 : }
254 20 : if (has_nan) {
255 0 : cfd_set_error(CFD_ERROR_DIVERGED,
256 : "NaN/Inf detected in energy_step_explicit_avx2");
257 0 : if (owns_buffer) cfd_free(T_new);
258 0 : return CFD_ERROR_DIVERGED;
259 : }
260 :
261 20 : memcpy(field->T, T_new, total * sizeof(double));
262 20 : if (owns_buffer) cfd_free(T_new);
263 :
264 : return CFD_SUCCESS;
265 : }
|
