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/memory.h"
13 :
14 : #include <math.h>
15 : #include <stdio.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 150649 : static inline int size_to_int(size_t val) {
50 150649 : if (val > (size_t)INT_MAX) {
51 : return INT_MAX;
52 : }
53 : return (int)val;
54 : }
55 :
56 64652 : 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 64652 : double sum = 0.0;
60 64652 : int ny_int = size_to_int(ny);
61 64652 : int nx_int = size_to_int(nx);
62 :
63 129454 : for (size_t k = k_start; k < k_end; k++) {
64 64802 : int j;
65 64802 : #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 64652 : return sum;
74 : }
75 :
76 42936 : 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 42936 : int ny_int = size_to_int(ny);
80 42936 : int nx_int = size_to_int(nx);
81 :
82 85928 : for (size_t k = k_start; k < k_end; k++) {
83 42992 : int j;
84 42992 : #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 42936 : }
93 :
94 21468 : 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 21468 : double dx2_inv = 1.0 / dx2;
99 21468 : double dy2_inv = 1.0 / dy2;
100 21468 : int ny_int = size_to_int(ny);
101 21468 : int nx_int = size_to_int(nx);
102 :
103 42964 : for (size_t k = k_start; k < k_end; k++) {
104 21496 : int j;
105 21496 : #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 21468 : }
118 :
119 124 : 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 124 : double dx2_inv = 1.0 / dx2;
124 124 : double dy2_inv = 1.0 / dy2;
125 124 : int ny_int = size_to_int(ny);
126 124 : int nx_int = size_to_int(nx);
127 :
128 281 : for (size_t k = k_start; k < k_end; k++) {
129 157 : int j;
130 157 : #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 124 : }
143 :
144 124 : 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 124 : int ny_int = size_to_int(ny);
148 :
149 281 : for (size_t k = k_start; k < k_end; k++) {
150 157 : int j;
151 157 : #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 21345 : 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 21345 : int ny_int = size_to_int(ny);
182 21345 : int nx_int = size_to_int(nx);
183 :
184 42690 : for (size_t k = k_start; k < k_end; k++) {
185 21345 : int j;
186 21345 : #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 21345 : }
195 :
196 : /* ============================================================================
197 : * CG OMP IMPLEMENTATION
198 : * ============================================================================ */
199 :
200 8 : 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 8 : cg_omp_context_t* ctx = (cg_omp_context_t*)cfd_calloc(1, sizeof(cg_omp_context_t));
207 8 : if (!ctx) {
208 : return CFD_ERROR_NOMEM;
209 : }
210 :
211 8 : ctx->dx2 = dx * dx;
212 8 : ctx->dy2 = dy * dy;
213 8 : ctx->inv_dz2 = poisson_solver_compute_inv_dz2(dz);
214 8 : poisson_solver_compute_3d_bounds(nz, nx, ny, &ctx->stride_z, &ctx->k_start, &ctx->k_end);
215 8 : ctx->diag_inv = 1.0 / (2.0 / ctx->dx2 + 2.0 / ctx->dy2 + 2.0 * ctx->inv_dz2);
216 8 : ctx->use_precond = (params && params->preconditioner == POISSON_PRECOND_JACOBI);
217 :
218 8 : size_t n = nx * ny * nz;
219 8 : ctx->r = (double*)cfd_calloc(n, sizeof(double));
220 8 : ctx->p = (double*)cfd_calloc(n, sizeof(double));
221 8 : ctx->Ap = (double*)cfd_calloc(n, sizeof(double));
222 8 : ctx->z = NULL;
223 :
224 8 : 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 8 : 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 8 : ctx->initialized = 1;
244 8 : solver->context = ctx;
245 8 : return CFD_SUCCESS;
246 : }
247 :
248 16 : static void cg_omp_destroy(poisson_solver_t* solver) {
249 16 : if (solver && solver->context) {
250 8 : cg_omp_context_t* ctx = (cg_omp_context_t*)solver->context;
251 8 : cfd_free(ctx->r);
252 8 : cfd_free(ctx->z);
253 8 : cfd_free(ctx->p);
254 8 : cfd_free(ctx->Ap);
255 8 : cfd_free(ctx);
256 8 : solver->context = NULL;
257 : }
258 16 : }
259 :
260 124 : 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 124 : (void)x_temp;
268 :
269 124 : cg_omp_context_t* ctx = (cg_omp_context_t*)solver->context;
270 124 : size_t nx = solver->nx;
271 124 : size_t ny = solver->ny;
272 124 : double dx2 = ctx->dx2;
273 124 : double dy2 = ctx->dy2;
274 124 : double inv_dz2 = ctx->inv_dz2;
275 124 : size_t k_start = ctx->k_start;
276 124 : size_t k_end = ctx->k_end;
277 124 : size_t stride_z = ctx->stride_z;
278 :
279 124 : double* r = ctx->r;
280 124 : double* z = ctx->z;
281 124 : double* p = ctx->p;
282 124 : double* Ap = ctx->Ap;
283 124 : int use_precond = ctx->use_precond;
284 124 : double diag_inv = ctx->diag_inv;
285 :
286 124 : poisson_solver_params_t* params = &solver->params;
287 124 : double start_time = poisson_solver_get_time_ms();
288 :
289 124 : poisson_solver_apply_bc(solver, x);
290 :
291 124 : compute_residual_omp(x, rhs, r, nx, ny, dx2, dy2, inv_dz2,
292 : k_start, k_end, stride_z);
293 :
294 124 : double rho;
295 124 : 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 124 : copy_vector_omp(r, p, nx, ny, k_start, k_end, stride_z);
302 124 : rho = dot_product_omp(r, r, nx, ny, k_start, k_end, stride_z);
303 : }
304 :
305 124 : double initial_res = sqrt(dot_product_omp(r, r, nx, ny,
306 : k_start, k_end, stride_z));
307 :
308 124 : if (stats) {
309 124 : stats->initial_residual = initial_res;
310 : }
311 :
312 124 : double tolerance = params->tolerance * initial_res;
313 124 : if (tolerance < params->absolute_tolerance) {
314 1 : tolerance = params->absolute_tolerance;
315 : }
316 :
317 124 : 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 21468 : int converged = 0;
328 : int iter;
329 : double res_norm = initial_res;
330 :
331 21468 : for (iter = 0; iter < params->max_iterations; iter++) {
332 21468 : apply_laplacian_omp(p, Ap, nx, ny, dx2, dy2, inv_dz2,
333 : k_start, k_end, stride_z);
334 :
335 21468 : double p_dot_Ap = dot_product_omp(p, Ap, nx, ny,
336 : k_start, k_end, stride_z);
337 21468 : CG_CHECK_BREAKDOWN(p_dot_Ap, stats, iter, res_norm, start_time);
338 :
339 21468 : double alpha = rho / p_dot_Ap;
340 :
341 21468 : axpy_omp(alpha, p, x, nx, ny, k_start, k_end, stride_z);
342 21468 : axpy_omp(-alpha, Ap, r, nx, ny, k_start, k_end, stride_z);
343 :
344 21468 : double rho_new;
345 21468 : 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 21468 : rho_new = dot_product_omp(r, r, nx, ny,
352 : k_start, k_end, stride_z);
353 : }
354 :
355 21468 : res_norm = sqrt(dot_product_omp(r, r, nx, ny,
356 : k_start, k_end, stride_z));
357 :
358 21468 : if (iter % params->check_interval == 0) {
359 21468 : if (params->verbose) {
360 0 : printf(" CG-OMP Iter %d: residual = %.6e\n", iter, res_norm);
361 : }
362 :
363 21468 : if (res_norm < tolerance || res_norm < params->absolute_tolerance) {
364 : converged = 1;
365 : break;
366 : }
367 : }
368 :
369 21345 : CG_CHECK_BREAKDOWN(rho, stats, iter, res_norm, start_time);
370 :
371 21345 : double beta = rho_new / rho;
372 42690 : update_search_direction_omp(use_precond ? z : r, p, beta, nx, ny,
373 : k_start, k_end, stride_z);
374 :
375 21345 : rho = rho_new;
376 : }
377 :
378 123 : if (!converged && (res_norm < tolerance || res_norm < params->absolute_tolerance)) {
379 123 : converged = 1;
380 : }
381 :
382 123 : poisson_solver_apply_bc(solver, x);
383 :
384 123 : double end_time = poisson_solver_get_time_ms();
385 :
386 123 : if (stats) {
387 123 : stats->iterations = (iter < params->max_iterations) ? (iter + 1) : iter;
388 123 : stats->final_residual = res_norm;
389 123 : stats->elapsed_time_ms = end_time - start_time;
390 123 : stats->status = converged ? POISSON_CONVERGED : POISSON_MAX_ITER;
391 : }
392 :
393 123 : 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 16 : poisson_solver_t* create_cg_omp_solver(void) {
416 16 : poisson_solver_t* solver = (poisson_solver_t*)cfd_calloc(1, sizeof(poisson_solver_t));
417 16 : if (!solver) {
418 : return NULL;
419 : }
420 :
421 16 : solver->name = POISSON_SOLVER_TYPE_CG_OMP;
422 16 : solver->description = "Conjugate Gradient (OpenMP)";
423 16 : solver->method = POISSON_METHOD_CG;
424 16 : solver->backend = POISSON_BACKEND_OMP;
425 16 : solver->params = poisson_solver_params_default();
426 :
427 16 : solver->init = cg_omp_init;
428 16 : solver->destroy = cg_omp_destroy;
429 16 : solver->solve = cg_omp_solve;
430 16 : solver->iterate = cg_omp_iterate;
431 16 : solver->apply_bc = NULL;
432 :
433 16 : return solver;
434 : }
435 :
436 : #endif /* CFD_ENABLE_OPENMP */
|
