Line data Source code
1 : /**
2 : * @file linear_solver_cg_omp.c
3 : * @brief Conjugate Gradient solver - OpenMP parallelized implementation
4 : *
5 : * Same algorithm as scalar CG but with OpenMP-parallelized primitives
6 : * (dot_product, axpy, apply_laplacian, etc.).
7 : */
8 :
9 : #include "../linear_solver_internal.h"
10 :
11 : #include "cfd/core/indexing.h"
12 : #include "cfd/core/logging.h"
13 : #include "cfd/core/memory.h"
14 :
15 : #include <math.h>
16 : #include <string.h>
17 :
18 : #ifdef CFD_ENABLE_OPENMP
19 :
20 : #include <omp.h>
21 :
22 : /* ============================================================================
23 : * CG CONTEXT
24 : * ============================================================================ */
25 :
26 : typedef struct {
27 : double dx2;
28 : double dy2;
29 : double inv_dz2; /* 1/dz^2 (0 for 2D) */
30 : double diag_inv;
31 :
32 : size_t stride_z; /* nx*ny for 3D, 0 for 2D */
33 : size_t k_start; /* first interior k index */
34 : size_t k_end; /* one-past-last interior k index */
35 :
36 : double* r;
37 : double* z;
38 : double* p;
39 : double* Ap;
40 :
41 : int use_precond;
42 : int initialized;
43 : } cg_omp_context_t;
44 :
45 : /* ============================================================================
46 : * OMP-PARALLELIZED PRIMITIVES
47 : * ============================================================================ */
48 :
49 153281 : static inline int size_to_int(size_t val) {
50 153281 : if (val > (size_t)INT_MAX) {
51 : return INT_MAX;
52 : }
53 : return (int)val;
54 : }
55 :
56 65795 : static double dot_product_omp(const double* a, const double* b,
57 : size_t nx, size_t ny,
58 : size_t k_start, size_t k_end, size_t stride_z) {
59 65795 : double sum = 0.0;
60 65795 : int ny_int = size_to_int(ny);
61 65795 : int nx_int = size_to_int(nx);
62 :
63 131740 : for (size_t k = k_start; k < k_end; k++) {
64 65945 : int j;
65 65945 : #pragma omp parallel for schedule(static) reduction(+:sum)
66 : for (j = 1; j < ny_int - 1; j++) {
67 : for (int i = 1; i < nx_int - 1; i++) {
68 : size_t idx = k * stride_z + IDX_2D((size_t)i, (size_t)j, nx);
69 : sum += a[idx] * b[idx];
70 : }
71 : }
72 : }
73 65795 : return sum;
74 : }
75 :
76 43670 : static void axpy_omp(double alpha, const double* x, double* y,
77 : size_t nx, size_t ny,
78 : size_t k_start, size_t k_end, size_t stride_z) {
79 43670 : int ny_int = size_to_int(ny);
80 43670 : int nx_int = size_to_int(nx);
81 :
82 87396 : for (size_t k = k_start; k < k_end; k++) {
83 43726 : int j;
84 43726 : #pragma omp parallel for schedule(static)
85 : for (j = 1; j < ny_int - 1; j++) {
86 : for (int i = 1; i < nx_int - 1; i++) {
87 : size_t idx = k * stride_z + IDX_2D((size_t)i, (size_t)j, nx);
88 : y[idx] += alpha * x[idx];
89 : }
90 : }
91 : }
92 43670 : }
93 :
94 21835 : static void apply_laplacian_omp(const double* p, double* Ap,
95 : size_t nx, size_t ny,
96 : double dx2, double dy2, double inv_dz2,
97 : size_t k_start, size_t k_end, size_t stride_z) {
98 21835 : double dx2_inv = 1.0 / dx2;
99 21835 : double dy2_inv = 1.0 / dy2;
100 21835 : int ny_int = size_to_int(ny);
101 21835 : int nx_int = size_to_int(nx);
102 :
103 43698 : for (size_t k = k_start; k < k_end; k++) {
104 21863 : int j;
105 21863 : #pragma omp parallel for schedule(static)
106 : for (j = 1; j < ny_int - 1; j++) {
107 : for (int i = 1; i < nx_int - 1; i++) {
108 : size_t idx = k * stride_z + IDX_2D((size_t)i, (size_t)j, nx);
109 : double laplacian =
110 : (p[idx + 1] - 2.0 * p[idx] + p[idx - 1]) * dx2_inv
111 : + (p[idx + nx] - 2.0 * p[idx] + p[idx - nx]) * dy2_inv
112 : + (p[idx + stride_z] + p[idx - stride_z] - 2.0 * p[idx]) * inv_dz2;
113 : Ap[idx] = -laplacian;
114 : }
115 : }
116 : }
117 21835 : }
118 :
119 145 : static void compute_residual_omp(const double* x, const double* rhs, double* r,
120 : size_t nx, size_t ny,
121 : double dx2, double dy2, double inv_dz2,
122 : size_t k_start, size_t k_end, size_t stride_z) {
123 145 : double dx2_inv = 1.0 / dx2;
124 145 : double dy2_inv = 1.0 / dy2;
125 145 : int ny_int = size_to_int(ny);
126 145 : int nx_int = size_to_int(nx);
127 :
128 323 : for (size_t k = k_start; k < k_end; k++) {
129 178 : int j;
130 178 : #pragma omp parallel for schedule(static)
131 : for (j = 1; j < ny_int - 1; j++) {
132 : for (int i = 1; i < nx_int - 1; i++) {
133 : size_t idx = k * stride_z + IDX_2D((size_t)i, (size_t)j, nx);
134 : double laplacian =
135 : (x[idx + 1] - 2.0 * x[idx] + x[idx - 1]) * dx2_inv
136 : + (x[idx + nx] - 2.0 * x[idx] + x[idx - nx]) * dy2_inv
137 : + (x[idx + stride_z] + x[idx - stride_z] - 2.0 * x[idx]) * inv_dz2;
138 : r[idx] = -rhs[idx] + laplacian;
139 : }
140 : }
141 : }
142 145 : }
143 :
144 145 : static void copy_vector_omp(const double* src, double* dst,
145 : size_t nx, size_t ny,
146 : size_t k_start, size_t k_end, size_t stride_z) {
147 145 : int ny_int = size_to_int(ny);
148 :
149 323 : for (size_t k = k_start; k < k_end; k++) {
150 178 : int j;
151 178 : #pragma omp parallel for schedule(static)
152 : for (j = 1; j < ny_int - 1; j++) {
153 : size_t row_start = k * stride_z + (size_t)j * nx;
154 : memcpy(&dst[row_start + 1], &src[row_start + 1], (nx - 2) * sizeof(double));
155 : }
156 : }
157 : }
158 :
159 0 : static void apply_jacobi_precond_omp(const double* r, double* z,
160 : size_t nx, size_t ny,
161 : double diag_inv,
162 : size_t k_start, size_t k_end, size_t stride_z) {
163 0 : int ny_int = size_to_int(ny);
164 0 : int nx_int = size_to_int(nx);
165 :
166 0 : for (size_t k = k_start; k < k_end; k++) {
167 0 : int j;
168 0 : #pragma omp parallel for schedule(static)
169 : for (j = 1; j < ny_int - 1; j++) {
170 : for (int i = 1; i < nx_int - 1; i++) {
171 : size_t idx = k * stride_z + IDX_2D((size_t)i, (size_t)j, nx);
172 : z[idx] = diag_inv * r[idx];
173 : }
174 : }
175 : }
176 0 : }
177 :
178 21691 : static void update_search_direction_omp(const double* src, double* p,
179 : double beta, size_t nx, size_t ny,
180 : size_t k_start, size_t k_end, size_t stride_z) {
181 21691 : int ny_int = size_to_int(ny);
182 21691 : int nx_int = size_to_int(nx);
183 :
184 43382 : for (size_t k = k_start; k < k_end; k++) {
185 21691 : int j;
186 21691 : #pragma omp parallel for schedule(static)
187 : for (j = 1; j < ny_int - 1; j++) {
188 : for (int i = 1; i < nx_int - 1; i++) {
189 : size_t idx = k * stride_z + IDX_2D((size_t)i, (size_t)j, nx);
190 : p[idx] = src[idx] + beta * p[idx];
191 : }
192 : }
193 : }
194 21691 : }
195 :
196 : /* ============================================================================
197 : * CG OMP IMPLEMENTATION
198 : * ============================================================================ */
199 :
200 10 : static cfd_status_t cg_omp_init(
201 : poisson_solver_t* solver,
202 : size_t nx, size_t ny, size_t nz,
203 : double dx, double dy, double dz,
204 : const poisson_solver_params_t* params)
205 : {
206 10 : cg_omp_context_t* ctx = (cg_omp_context_t*)cfd_calloc(1, sizeof(cg_omp_context_t));
207 10 : if (!ctx) {
208 : return CFD_ERROR_NOMEM;
209 : }
210 :
211 10 : ctx->dx2 = dx * dx;
212 10 : ctx->dy2 = dy * dy;
213 10 : ctx->inv_dz2 = poisson_solver_compute_inv_dz2(dz);
214 10 : poisson_solver_compute_3d_bounds(nz, nx, ny, &ctx->stride_z, &ctx->k_start, &ctx->k_end);
215 10 : ctx->diag_inv = 1.0 / (2.0 / ctx->dx2 + 2.0 / ctx->dy2 + 2.0 * ctx->inv_dz2);
216 10 : ctx->use_precond = (params && params->preconditioner == POISSON_PRECOND_JACOBI);
217 :
218 10 : size_t n = nx * ny * nz;
219 10 : ctx->r = (double*)cfd_calloc(n, sizeof(double));
220 10 : ctx->p = (double*)cfd_calloc(n, sizeof(double));
221 10 : ctx->Ap = (double*)cfd_calloc(n, sizeof(double));
222 10 : ctx->z = NULL;
223 :
224 10 : if (!ctx->r || !ctx->p || !ctx->Ap) {
225 0 : cfd_free(ctx->r);
226 0 : cfd_free(ctx->p);
227 0 : cfd_free(ctx->Ap);
228 0 : cfd_free(ctx);
229 0 : return CFD_ERROR_NOMEM;
230 : }
231 :
232 10 : if (ctx->use_precond) {
233 0 : ctx->z = (double*)cfd_calloc(n, sizeof(double));
234 0 : if (!ctx->z) {
235 0 : cfd_free(ctx->r);
236 0 : cfd_free(ctx->p);
237 0 : cfd_free(ctx->Ap);
238 0 : cfd_free(ctx);
239 0 : return CFD_ERROR_NOMEM;
240 : }
241 : }
242 :
243 10 : ctx->initialized = 1;
244 10 : solver->context = ctx;
245 10 : return CFD_SUCCESS;
246 : }
247 :
248 19 : static void cg_omp_destroy(poisson_solver_t* solver) {
249 19 : if (solver && solver->context) {
250 10 : cg_omp_context_t* ctx = (cg_omp_context_t*)solver->context;
251 10 : cfd_free(ctx->r);
252 10 : cfd_free(ctx->z);
253 10 : cfd_free(ctx->p);
254 10 : cfd_free(ctx->Ap);
255 10 : cfd_free(ctx);
256 10 : solver->context = NULL;
257 : }
258 19 : }
259 :
260 145 : static cfd_status_t cg_omp_solve(
261 : poisson_solver_t* solver,
262 : double* x,
263 : double* x_temp,
264 : const double* rhs,
265 : poisson_solver_stats_t* stats)
266 : {
267 145 : (void)x_temp;
268 :
269 145 : cg_omp_context_t* ctx = (cg_omp_context_t*)solver->context;
270 145 : size_t nx = solver->nx;
271 145 : size_t ny = solver->ny;
272 145 : double dx2 = ctx->dx2;
273 145 : double dy2 = ctx->dy2;
274 145 : double inv_dz2 = ctx->inv_dz2;
275 145 : size_t k_start = ctx->k_start;
276 145 : size_t k_end = ctx->k_end;
277 145 : size_t stride_z = ctx->stride_z;
278 :
279 145 : double* r = ctx->r;
280 145 : double* z = ctx->z;
281 145 : double* p = ctx->p;
282 145 : double* Ap = ctx->Ap;
283 145 : int use_precond = ctx->use_precond;
284 145 : double diag_inv = ctx->diag_inv;
285 :
286 145 : poisson_solver_params_t* params = &solver->params;
287 145 : double start_time = poisson_solver_get_time_ms();
288 :
289 145 : poisson_solver_apply_bc(solver, x);
290 :
291 145 : compute_residual_omp(x, rhs, r, nx, ny, dx2, dy2, inv_dz2,
292 : k_start, k_end, stride_z);
293 :
294 145 : double rho;
295 145 : if (use_precond) {
296 0 : apply_jacobi_precond_omp(r, z, nx, ny, diag_inv,
297 : k_start, k_end, stride_z);
298 0 : copy_vector_omp(z, p, nx, ny, k_start, k_end, stride_z);
299 0 : rho = dot_product_omp(r, z, nx, ny, k_start, k_end, stride_z);
300 : } else {
301 145 : copy_vector_omp(r, p, nx, ny, k_start, k_end, stride_z);
302 145 : rho = dot_product_omp(r, r, nx, ny, k_start, k_end, stride_z);
303 : }
304 :
305 145 : double initial_res = sqrt(dot_product_omp(r, r, nx, ny,
306 : k_start, k_end, stride_z));
307 :
308 145 : if (stats) {
309 145 : stats->initial_residual = initial_res;
310 : }
311 :
312 145 : double tolerance = params->tolerance * initial_res;
313 145 : if (tolerance < params->absolute_tolerance) {
314 1 : tolerance = params->absolute_tolerance;
315 : }
316 :
317 145 : if (initial_res < params->absolute_tolerance) {
318 1 : if (stats) {
319 1 : stats->status = POISSON_CONVERGED;
320 1 : stats->iterations = 0;
321 1 : stats->final_residual = initial_res;
322 1 : stats->elapsed_time_ms = poisson_solver_get_time_ms() - start_time;
323 : }
324 1 : return CFD_SUCCESS;
325 : }
326 :
327 21835 : int converged = 0;
328 : int iter;
329 : double res_norm = initial_res;
330 :
331 21835 : for (iter = 0; iter < params->max_iterations; iter++) {
332 21835 : apply_laplacian_omp(p, Ap, nx, ny, dx2, dy2, inv_dz2,
333 : k_start, k_end, stride_z);
334 :
335 21835 : double p_dot_Ap = dot_product_omp(p, Ap, nx, ny,
336 : k_start, k_end, stride_z);
337 21835 : CG_CHECK_BREAKDOWN(p_dot_Ap, stats, iter, res_norm, start_time);
338 :
339 21835 : double alpha = rho / p_dot_Ap;
340 :
341 21835 : axpy_omp(alpha, p, x, nx, ny, k_start, k_end, stride_z);
342 21835 : axpy_omp(-alpha, Ap, r, nx, ny, k_start, k_end, stride_z);
343 :
344 21835 : double rho_new;
345 21835 : if (use_precond) {
346 0 : apply_jacobi_precond_omp(r, z, nx, ny, diag_inv,
347 : k_start, k_end, stride_z);
348 0 : rho_new = dot_product_omp(r, z, nx, ny,
349 : k_start, k_end, stride_z);
350 : } else {
351 21835 : rho_new = dot_product_omp(r, r, nx, ny,
352 : k_start, k_end, stride_z);
353 : }
354 :
355 21835 : res_norm = sqrt(dot_product_omp(r, r, nx, ny,
356 : k_start, k_end, stride_z));
357 :
358 21835 : if (iter % params->check_interval == 0) {
359 21835 : if (params->verbose) {
360 0 : CFD_LOG_DEBUG("poisson", "CG-OMP Iter %d: residual = %.6e", iter, res_norm);
361 : }
362 :
363 21835 : if (res_norm < tolerance || res_norm < params->absolute_tolerance) {
364 : converged = 1;
365 : break;
366 : }
367 : }
368 :
369 21691 : CG_CHECK_BREAKDOWN(rho, stats, iter, res_norm, start_time);
370 :
371 21691 : double beta = rho_new / rho;
372 43382 : update_search_direction_omp(use_precond ? z : r, p, beta, nx, ny,
373 : k_start, k_end, stride_z);
374 :
375 21691 : rho = rho_new;
376 : }
377 :
378 144 : if (!converged && (res_norm < tolerance || res_norm < params->absolute_tolerance)) {
379 144 : converged = 1;
380 : }
381 :
382 144 : poisson_solver_apply_bc(solver, x);
383 :
384 144 : double end_time = poisson_solver_get_time_ms();
385 :
386 144 : if (stats) {
387 144 : stats->iterations = (iter < params->max_iterations) ? (iter + 1) : iter;
388 144 : stats->final_residual = res_norm;
389 144 : stats->elapsed_time_ms = end_time - start_time;
390 144 : stats->status = converged ? POISSON_CONVERGED : POISSON_MAX_ITER;
391 : }
392 :
393 144 : return converged ? CFD_SUCCESS : CFD_ERROR_MAX_ITER;
394 : }
395 :
396 0 : static cfd_status_t cg_omp_iterate(
397 : poisson_solver_t* solver,
398 : double* x,
399 : double* x_temp,
400 : const double* rhs,
401 : double* residual)
402 : {
403 0 : (void)x_temp;
404 :
405 0 : if (residual) {
406 0 : *residual = poisson_solver_compute_residual(solver, x, rhs);
407 : }
408 0 : return CFD_SUCCESS;
409 : }
410 :
411 : /* ============================================================================
412 : * FACTORY FUNCTION
413 : * ============================================================================ */
414 :
415 19 : poisson_solver_t* create_cg_omp_solver(void) {
416 19 : poisson_solver_t* solver = (poisson_solver_t*)cfd_calloc(1, sizeof(poisson_solver_t));
417 19 : if (!solver) {
418 : return NULL;
419 : }
420 :
421 19 : solver->name = POISSON_SOLVER_TYPE_CG_OMP;
422 19 : solver->description = "Conjugate Gradient (OpenMP)";
423 19 : solver->method = POISSON_METHOD_CG;
424 19 : solver->backend = POISSON_BACKEND_OMP;
425 19 : solver->params = poisson_solver_params_default();
426 :
427 19 : solver->init = cg_omp_init;
428 19 : solver->destroy = cg_omp_destroy;
429 19 : solver->solve = cg_omp_solve;
430 19 : solver->iterate = cg_omp_iterate;
431 19 : solver->apply_bc = NULL;
432 :
433 19 : return solver;
434 : }
435 :
436 : #endif /* CFD_ENABLE_OPENMP */
|
