Loading...
Searching...
No Matches
grid_operations.hpp
Go to the documentation of this file.
1#pragma once
2
3#include "grid/grid_types.hpp"
4#include "kokkos/kokkos_wrapper.hpp"
5#include "mpi/mpi.hpp"
6#include "terra/grid/bit_masks.hpp"
7#include "terra/util/bit_masking.hpp"
8
9namespace terra::kernels::common {
10
11template < typename ScalarType >
12void set_constant( const grid::Grid2DDataScalar< ScalarType >& x, ScalarType value )
13{
14 Kokkos::deep_copy( x, value );
15}
16
17template < typename ScalarType >
18void set_constant( const grid::Grid3DDataScalar< ScalarType >& x, ScalarType value )
19{
20 Kokkos::deep_copy( x, value );
21}
22
23template < typename ScalarType >
24void set_constant( const grid::Grid4DDataScalar< ScalarType >& x, ScalarType value )
25{
26 Kokkos::deep_copy( x, value );
27}
28
29template < typename ScalarType, int VecDim >
30void set_constant( const grid::Grid4DDataVec< ScalarType, VecDim >& x, ScalarType value )
31{
32 for ( int d = 0; d < VecDim; ++d )
33 set_constant( x.comp_[d], value );
34}
35
36template < typename ScalarType >
37void set_constant( const grid::Grid5DDataScalar< ScalarType >& x, ScalarType value )
38{
39 Kokkos::deep_copy( x, value );
40}
41
42template < typename ScalarType >
43void scale( const grid::Grid4DDataScalar< ScalarType >& x, ScalarType value )
44{
45 Kokkos::parallel_for(
46 "scale (Grid3DDataScalar)",
47 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
48 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) { x( local_subdomain, i, j, k ) *= value; } );
49
50 Kokkos::fence();
51}
52
53template < typename ScalarType, util::FlagLike FlagType >
54void assign_masked_else_keep_old(
55 const grid::Grid4DDataScalar< ScalarType >& dst,
56 const ScalarType& value,
57 const grid::Grid4DDataScalar< FlagType >& mask_grid,
58 const FlagType mask_value )
59{
60 Kokkos::parallel_for(
61 "assign_masked",
62 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { dst.extent( 0 ), dst.extent( 1 ), dst.extent( 2 ), dst.extent( 3 ) } ),
63 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
64 const ScalarType mask_val = util::has_flag( mask_grid( local_subdomain, i, j, k ), mask_value ) ? 1.0 : 0.0;
65 dst( local_subdomain, i, j, k ) = mask_val * value + ( 1.0 - mask_val ) * dst( local_subdomain, i, j, k );
66 } );
67
68 Kokkos::fence();
69}
70
71template < typename ScalarType, util::FlagLike FlagType >
72void assign_masked_else_keep_old(
73 const grid::Grid4DDataScalar< ScalarType >& dst,
74 const grid::Grid4DDataScalar< ScalarType >& src,
75 const grid::Grid4DDataScalar< FlagType >& mask_grid,
76 const FlagType mask_value )
77{
78 Kokkos::parallel_for(
79 "assign_masked",
80 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { dst.extent( 0 ), dst.extent( 1 ), dst.extent( 2 ), dst.extent( 3 ) } ),
81 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
82 const ScalarType mask_val = util::has_flag( mask_grid( local_subdomain, i, j, k ), mask_value ) ? 1.0 : 0.0;
83 dst( local_subdomain, i, j, k ) =
84 mask_val * src( local_subdomain, i, j, k ) + ( 1.0 - mask_val ) * dst( local_subdomain, i, j, k );
85 } );
86
87 Kokkos::fence();
88}
89
90template < typename ScalarType, int VecDim, util::FlagLike FlagType >
91void assign_masked_else_keep_old(
92 const grid::Grid4DDataVec< ScalarType, VecDim >& dst,
93 const ScalarType& value,
94 const grid::Grid4DDataScalar< FlagType >& mask_grid,
95 const FlagType mask_value )
96{
97 for ( int d = 0; d < VecDim; ++d )
98 assign_masked_else_keep_old( dst.comp_[d], value, mask_grid, mask_value );
99}
100
101template < typename ScalarType, int VecDim, util::FlagLike FlagType >
102void assign_masked_else_keep_old(
103 const grid::Grid4DDataVec< ScalarType, VecDim >& dst,
104 const grid::Grid4DDataVec< ScalarType, VecDim >& src,
105 const grid::Grid4DDataScalar< FlagType >& mask_grid,
106 const FlagType mask_value )
107{
108 for ( int d = 0; d < VecDim; ++d )
109 assign_masked_else_keep_old( dst.comp_[d], src.comp_[d], mask_grid, mask_value );
110}
111
112template < typename ScalarType, int VecDim, util::FlagLike FlagType >
113void assign_masked_else_keep_old(
114 const grid::Grid4DDataVec< ScalarType, VecDim >& dst,
115 const ScalarType& value,
116 const grid::Grid4DDataScalar< FlagType >& mask_grid,
117 const FlagType mask_value,
118 const int vector_component )
119{
120 Kokkos::parallel_for(
121 "assign_masked",
122 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { dst.extent( 0 ), dst.extent( 1 ), dst.extent( 2 ), dst.extent( 3 ) } ),
123 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
124 const ScalarType mask_val = util::has_flag( mask_grid( local_subdomain, i, j, k ), mask_value ) ? 1.0 : 0.0;
125 dst( local_subdomain, i, j, k, vector_component ) =
126 mask_val * value + ( 1.0 - mask_val ) * dst( local_subdomain, i, j, k, vector_component );
127 } );
128
129 Kokkos::fence();
130}
131
132template < typename ScalarType >
133void lincomb(
134 const grid::Grid4DDataScalar< ScalarType >& y,
135 ScalarType c_0,
136 ScalarType c_1,
137 const grid::Grid4DDataScalar< ScalarType >& x_1 )
138{
139 Kokkos::parallel_for(
140 "lincomb 1 arg (Grid4DDataScalar)",
141 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { y.extent( 0 ), y.extent( 1 ), y.extent( 2 ), y.extent( 3 ) } ),
142 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
143 y( local_subdomain, i, j, k ) = c_0 + c_1 * x_1( local_subdomain, i, j, k );
144 } );
145
146 Kokkos::fence();
147}
148
149template < typename ScalarType >
150void lincomb(
151 const grid::Grid4DDataScalar< ScalarType >& y,
152 ScalarType c_0,
153 ScalarType c_1,
154 const grid::Grid4DDataScalar< ScalarType >& x_1,
155 ScalarType c_2,
156 const grid::Grid4DDataScalar< ScalarType >& x_2 )
157{
158 Kokkos::parallel_for(
159 "lincomb 2 args (Grid4DDataScalar)",
160 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { y.extent( 0 ), y.extent( 1 ), y.extent( 2 ), y.extent( 3 ) } ),
161 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
162 y( local_subdomain, i, j, k ) =
163 c_0 + c_1 * x_1( local_subdomain, i, j, k ) + c_2 * x_2( local_subdomain, i, j, k );
164 } );
165
166 Kokkos::fence();
167}
168
169template < typename ScalarType >
170void lincomb(
171 const grid::Grid4DDataScalar< ScalarType >& y,
172 ScalarType c_0,
173 ScalarType c_1,
174 const grid::Grid4DDataScalar< ScalarType >& x_1,
175 ScalarType c_2,
176 const grid::Grid4DDataScalar< ScalarType >& x_2,
177 ScalarType c_3,
178 const grid::Grid4DDataScalar< ScalarType >& x_3 )
179{
180 Kokkos::parallel_for(
181 "lincomb 3 args (Grid4DDataScalar)",
182 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { y.extent( 0 ), y.extent( 1 ), y.extent( 2 ), y.extent( 3 ) } ),
183 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
184 y( local_subdomain, i, j, k ) = c_0 + c_1 * x_1( local_subdomain, i, j, k ) +
185 c_2 * x_2( local_subdomain, i, j, k ) +
186 c_3 * x_3( local_subdomain, i, j, k );
187 } );
188
189 Kokkos::fence();
190}
191
192template < typename ScalarType, int VecDim >
193void lincomb(
194 const grid::Grid4DDataVec< ScalarType, VecDim >& y,
195 ScalarType c_0,
196 ScalarType c_1,
197 const grid::Grid4DDataVec< ScalarType, VecDim >& x_1 )
198{
199 for ( int d = 0; d < VecDim; ++d )
200 lincomb( y.comp_[d], c_0, c_1, x_1.comp_[d] );
201}
202
203template < typename ScalarType, int VecDim >
204void lincomb(
205 const grid::Grid4DDataVec< ScalarType, VecDim >& y,
206 ScalarType c_0,
207 ScalarType c_1,
208 const grid::Grid4DDataVec< ScalarType, VecDim >& x_1,
209 ScalarType c_2,
210 const grid::Grid4DDataVec< ScalarType, VecDim >& x_2 )
211{
212 for ( int d = 0; d < VecDim; ++d )
213 lincomb( y.comp_[d], c_0, c_1, x_1.comp_[d], c_2, x_2.comp_[d] );
214}
215
216template < typename ScalarType, int VecDim >
217void lincomb(
218 const grid::Grid4DDataVec< ScalarType, VecDim >& y,
219 ScalarType c_0,
220 ScalarType c_1,
221 const grid::Grid4DDataVec< ScalarType, VecDim >& x_1,
222 ScalarType c_2,
223 const grid::Grid4DDataVec< ScalarType, VecDim >& x_2,
224 ScalarType c_3,
225 const grid::Grid4DDataVec< ScalarType, VecDim >& x_3 )
226{
227 for ( int d = 0; d < VecDim; ++d )
228 lincomb( y.comp_[d], c_0, c_1, x_1.comp_[d], c_2, x_2.comp_[d], c_3, x_3.comp_[d] );
229}
230
231template < typename ScalarType >
232void invert_inplace( const grid::Grid4DDataScalar< ScalarType >& y )
233{
234 Kokkos::parallel_for(
235 "invert",
236 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { y.extent( 0 ), y.extent( 1 ), y.extent( 2 ), y.extent( 3 ) } ),
237 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
238 y( local_subdomain, i, j, k ) = 1.0 / y( local_subdomain, i, j, k );
239 } );
240
241 Kokkos::fence();
242}
243
244template < typename ScalarType, int VecDim >
245void invert_inplace( const grid::Grid4DDataVec< ScalarType, VecDim >& y )
246{
247 for ( int d = 0; d < VecDim; ++d )
248 invert_inplace( y.comp_[d] );
249}
250
251template < typename ScalarType >
252void mult_elementwise_inplace(
253 const grid::Grid4DDataScalar< ScalarType >& y,
254 const grid::Grid4DDataScalar< ScalarType >& x )
255{
256 Kokkos::parallel_for(
257 "mult_elementwise_inplace",
258 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { y.extent( 0 ), y.extent( 1 ), y.extent( 2 ), y.extent( 3 ) } ),
259 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
260 y( local_subdomain, i, j, k ) *= x( local_subdomain, i, j, k );
261 } );
262
263 Kokkos::fence();
264}
265
266template < typename ScalarType, int VecDim >
267void mult_elementwise_inplace(
268 const grid::Grid4DDataVec< ScalarType, VecDim >& y,
269 const grid::Grid4DDataVec< ScalarType, VecDim >& x )
270{
271 for ( int d = 0; d < VecDim; ++d )
272 mult_elementwise_inplace( y.comp_[d], x.comp_[d] );
273}
274
275template < typename ScalarType >
276ScalarType min_entry( const grid::Grid4DDataScalar< ScalarType >& x, MPI_Comm comm = MPI_COMM_WORLD )
277{
278 ScalarType min_val = 0.0;
279 Kokkos::parallel_reduce(
280 "min_entry",
281 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
282 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_min ) {
283 ScalarType val = x( local_subdomain, i, j, k );
284 local_min = Kokkos::min( local_min, val );
285 },
286 Kokkos::Min< ScalarType >( min_val ) );
287
288 Kokkos::fence();
289
290 MPI_Allreduce( MPI_IN_PLACE, &min_val, 1, mpi::mpi_datatype< ScalarType >(), MPI_MIN, comm );
291
292 return min_val;
293}
294
295template < typename ScalarType >
296ScalarType min_abs_entry( const grid::Grid4DDataScalar< ScalarType >& x, MPI_Comm comm = MPI_COMM_WORLD )
297{
298 ScalarType min_mag = 0.0;
299 Kokkos::parallel_reduce(
300 "min_abs_entry",
301 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
302 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_min ) {
303 ScalarType val = Kokkos::abs( x( local_subdomain, i, j, k ) );
304 local_min = Kokkos::min( local_min, val );
305 },
306 Kokkos::Min< ScalarType >( min_mag ) );
307
308 Kokkos::fence();
309
310 MPI_Allreduce( MPI_IN_PLACE, &min_mag, 1, mpi::mpi_datatype< ScalarType >(), MPI_MIN, comm );
311
312 return min_mag;
313}
314
315template < typename ScalarType >
316ScalarType max_abs_entry( const grid::Grid4DDataScalar< ScalarType >& x, MPI_Comm comm = MPI_COMM_WORLD )
317{
318 ScalarType max_mag = 0.0;
319 Kokkos::parallel_reduce(
320 "max_abs_entry",
321 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
322 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_max ) {
323 ScalarType val = Kokkos::abs( x( local_subdomain, i, j, k ) );
324 local_max = Kokkos::max( local_max, val );
325 },
326 Kokkos::Max< ScalarType >( max_mag ) );
327
328 Kokkos::fence();
329
330 MPI_Allreduce( MPI_IN_PLACE, &max_mag, 1, mpi::mpi_datatype< ScalarType >(), MPI_MAX, comm );
331
332 return max_mag;
333}
334
335template < typename ScalarType, int VecDim >
336ScalarType max_abs_entry( const grid::Grid4DDataVec< ScalarType, VecDim >& x, MPI_Comm comm = MPI_COMM_WORLD )
337{
338 ScalarType max_mag = 0.0;
339 for ( int d = 0; d < VecDim; ++d )
340 {
341 // max_abs_entry for scalar internally does MPI_Allreduce,
342 // so we call per-component and take the max on the host.
343 ScalarType comp_max = max_abs_entry( x.comp_[d], comm );
344 max_mag = std::max( max_mag, comp_max );
345 }
346 return max_mag;
347}
348
349template < typename ScalarType, util::FlagLike FlagType >
350ScalarType max_abs_entry(
351 const grid::Grid4DDataScalar< ScalarType >& x,
352 const grid::Grid4DDataScalar< FlagType >& mask,
353 const FlagType& mask_value,
354 MPI_Comm comm = MPI_COMM_WORLD )
355{
356 ScalarType max_mag = 0.0;
357 Kokkos::parallel_reduce(
358 "max_abs_entry",
359 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
360 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_max ) {
361 if ( util::has_flag( mask( local_subdomain, i, j, k ), mask_value ) )
362 {
363 ScalarType val = Kokkos::abs( x( local_subdomain, i, j, k ) );
364 local_max = Kokkos::max( local_max, val );
365 }
366 },
367 Kokkos::Max< ScalarType >( max_mag ) );
368
369 Kokkos::fence();
370
371 MPI_Allreduce( MPI_IN_PLACE, &max_mag, 1, mpi::mpi_datatype< ScalarType >(), MPI_MAX, comm );
372
373 return max_mag;
374}
375
376template < typename ScalarType >
377ScalarType max_abs_entry_subset(
378 const grid::Grid4DDataScalar< ScalarType >& x,
379 dense::Vec< int, 4 > start,
380 dense::Vec< int, 4 > end_excl,
381 MPI_Comm comm = MPI_COMM_WORLD )
382{
383 ScalarType max_mag = 0.0;
384 Kokkos::parallel_reduce(
385 "max_abs_entry",
386 Kokkos::MDRangePolicy(
387 { start( 0 ), start( 1 ), start( 2 ), start( 3 ) },
388 { end_excl( 0 ), end_excl( 1 ), end_excl( 2 ), end_excl( 3 ) } ),
389 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_max ) {
390 ScalarType val = Kokkos::abs( x( local_subdomain, i, j, k ) );
391 local_max = Kokkos::max( local_max, val );
392 },
393 Kokkos::Max< ScalarType >( max_mag ) );
394
395 Kokkos::fence();
396
397 MPI_Allreduce( MPI_IN_PLACE, &max_mag, 1, mpi::mpi_datatype< ScalarType >(), MPI_MAX, comm );
398
399 return max_mag;
400}
401
402template < typename ScalarType, int VecDim >
403ScalarType max_vector_magnitude( const grid::Grid4DDataVec< ScalarType, VecDim >& x, MPI_Comm comm = MPI_COMM_WORLD )
404{
405 ScalarType max_mag = 0.0;
406 Kokkos::parallel_reduce(
407 "max_vector_magnitude",
408 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
409 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_max ) {
410 ScalarType val = 0;
411 for ( int d = 0; d < VecDim; ++d )
412 {
413 val += x( local_subdomain, i, j, k, d ) * x( local_subdomain, i, j, k, d );
414 }
415 val = Kokkos::sqrt( val );
416 local_max = Kokkos::max( local_max, val );
417 },
418 Kokkos::Max< ScalarType >( max_mag ) );
419
420 Kokkos::fence();
421
422 MPI_Allreduce( MPI_IN_PLACE, &max_mag, 1, mpi::mpi_datatype< ScalarType >(), MPI_MAX, comm );
423
424 return max_mag;
425}
426
427template < typename ScalarType, int VecDim >
428void vector_magnitude(
429 grid::Grid4DDataScalar< ScalarType >& magnitude_out,
430 const grid::Grid4DDataVec< ScalarType, VecDim >& vectorial_data_in )
431{
432 Kokkos::parallel_for(
433 "vector_magnitude",
434 Kokkos::MDRangePolicy(
435 { 0, 0, 0, 0 },
436 { vectorial_data_in.extent( 0 ),
437 vectorial_data_in.extent( 1 ),
438 vectorial_data_in.extent( 2 ),
439 vectorial_data_in.extent( 3 ) } ),
440 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
441 ScalarType val = 0;
442 for ( int d = 0; d < VecDim; ++d )
443 {
444 val +=
445 vectorial_data_in( local_subdomain, i, j, k, d ) * vectorial_data_in( local_subdomain, i, j, k, d );
446 }
447 magnitude_out( local_subdomain, i, j, k ) = Kokkos::sqrt( val );
448 } );
449
450 Kokkos::fence();
451}
452
453template < typename ScalarType, int VecDim >
454void extract_vector_component(
455 grid::Grid4DDataScalar< ScalarType >& component_out,
456 const grid::Grid4DDataVec< ScalarType, VecDim >& vectorial_data_in,
457 const int component )
458{
459 if ( component < 0 || component >= VecDim )
460 {
461 Kokkos::abort( "Vector component invalid." );
462 }
463
464 Kokkos::parallel_for(
465 "extract_vector_component",
466 Kokkos::MDRangePolicy(
467 { 0, 0, 0, 0 },
468 { vectorial_data_in.extent( 0 ),
469 vectorial_data_in.extent( 1 ),
470 vectorial_data_in.extent( 2 ),
471 vectorial_data_in.extent( 3 ) } ),
472 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
473 component_out( local_subdomain, i, j, k ) = vectorial_data_in( local_subdomain, i, j, k, component );
474 } );
475
476 Kokkos::fence();
477}
478
479template < typename ScalarType, int VecDim >
480void set_vector_component(
481 grid::Grid4DDataVec< ScalarType, VecDim >& vectorial_data,
482 const int component,
483 const ScalarType constant )
484{
485 if ( component < 0 || component >= VecDim )
486 {
487 Kokkos::abort( "Vector component invalid." );
488 }
489
490 Kokkos::parallel_for(
491 "set_vector_component",
492 Kokkos::MDRangePolicy(
493 { 0, 0, 0, 0 },
494 { vectorial_data.extent( 0 ),
495 vectorial_data.extent( 1 ),
496 vectorial_data.extent( 2 ),
497 vectorial_data.extent( 3 ) } ),
498 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
499 vectorial_data( local_subdomain, i, j, k, component ) = constant;
500 } );
501
502 Kokkos::fence();
503}
504
505template < typename ScalarType >
506ScalarType sum_of_absolutes( const grid::Grid4DDataScalar< ScalarType >& x, MPI_Comm comm = MPI_COMM_WORLD )
507{
508 ScalarType sum_abs = 0.0;
509 Kokkos::parallel_reduce(
510 "sum_of_absolutes",
511 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
512 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_sum_abs ) {
513 ScalarType val = Kokkos::abs( x( local_subdomain, i, j, k ) );
514 local_sum_abs = local_sum_abs + val;
515 },
516 Kokkos::Sum< ScalarType >( sum_abs ) );
517
518 Kokkos::fence();
519
520 MPI_Allreduce( MPI_IN_PLACE, &sum_abs, 1, mpi::mpi_datatype< ScalarType >(), MPI_SUM, comm );
521
522 return sum_abs;
523}
524
525template < typename ScalarType, util::FlagLike FlagType >
526ScalarType count_masked( const grid::Grid4DDataScalar< FlagType >& mask,
527 const FlagType& mask_value,
528 MPI_Comm comm = MPI_COMM_WORLD )
529{
530 auto count = static_cast< ScalarType >( 0 );
531
532 Kokkos::parallel_reduce(
533 "masked_sum",
534 Kokkos::MDRangePolicy(
535 { 0, 0, 0, 0 }, { mask.extent( 0 ), mask.extent( 1 ), mask.extent( 2 ), mask.extent( 3 ) } ),
536 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_sum ) {
537 const ScalarType mask_val = util::has_flag( mask( local_subdomain, i, j, k ), mask_value ) ?
538 static_cast< ScalarType >( 1 ) :
539 static_cast< ScalarType >( 0 );
540 local_sum = local_sum + mask_val;
541 },
542 Kokkos::Sum< ScalarType >( count ) );
543
544 Kokkos::fence();
545
546 MPI_Allreduce( MPI_IN_PLACE, &count, 1, mpi::mpi_datatype< ScalarType >(), MPI_SUM, comm );
547
548 return count;
549}
550
551template < typename ScalarType, util::FlagLike FlagType >
552ScalarType masked_sum(
553 const grid::Grid4DDataScalar< ScalarType >& x,
554 const grid::Grid4DDataScalar< FlagType >& mask,
555 const FlagType& mask_value,
556 MPI_Comm comm = MPI_COMM_WORLD )
557{
558 ScalarType sum = 0.0;
559
560 Kokkos::parallel_reduce(
561 "masked_sum",
562 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
563 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_sum ) {
564 const ScalarType mask_val = util::has_flag( mask( local_subdomain, i, j, k ), mask_value ) ? 1.0 : 0.0;
565 ScalarType val = x( local_subdomain, i, j, k ) * mask_val;
566 local_sum = local_sum + val;
567 },
568 Kokkos::Sum< ScalarType >( sum ) );
569
570 Kokkos::fence();
571
572 MPI_Allreduce( MPI_IN_PLACE, &sum, 1, mpi::mpi_datatype< ScalarType >(), MPI_SUM, comm );
573
574 return sum;
575}
576
577template < typename ScalarType, util::FlagLike FlagType0, util::FlagLike FlagType1 >
578ScalarType masked_sum(
579 const grid::Grid4DDataScalar< ScalarType >& x,
580 const grid::Grid4DDataScalar< FlagType0 >& mask0,
581 const grid::Grid4DDataScalar< FlagType1 >& mask1,
582 const FlagType0& mask0_value,
583 const FlagType1& mask1_value,
584 MPI_Comm comm = MPI_COMM_WORLD )
585{
586 ScalarType sum = 0.0;
587
588 Kokkos::parallel_reduce(
589 "masked_sum",
590 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
591 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_sum ) {
592 ScalarType mask_val = 1.0;
593 mask_val *= util::has_flag( mask0( local_subdomain, i, j, k ), mask0_value ) ? 1.0 : 0.0;
594 mask_val *= util::has_flag( mask1( local_subdomain, i, j, k ), mask1_value ) ? 1.0 : 0.0;
595 ScalarType val = x( local_subdomain, i, j, k ) * mask_val;
596 local_sum = local_sum + val;
597 },
598 Kokkos::Sum< ScalarType >( sum ) );
599
600 Kokkos::fence();
601
602 MPI_Allreduce( MPI_IN_PLACE, &sum, 1, mpi::mpi_datatype< ScalarType >(), MPI_SUM, comm );
603
604 return sum;
605}
606
607template < typename ScalarType >
608ScalarType dot_product( const grid::Grid4DDataScalar< ScalarType >& x,
609 const grid::Grid4DDataScalar< ScalarType >& y,
610 MPI_Comm comm = MPI_COMM_WORLD )
611{
612 ScalarType dot_prod = 0.0;
613
614 Kokkos::parallel_reduce(
615 "dot_product",
616 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
617 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_dot_prod ) {
618 ScalarType val = x( local_subdomain, i, j, k ) * y( local_subdomain, i, j, k );
619 local_dot_prod = local_dot_prod + val;
620 },
621 Kokkos::Sum< ScalarType >( dot_prod ) );
622
623 Kokkos::fence( "dot_product" );
624
625 MPI_Allreduce( MPI_IN_PLACE, &dot_prod, 1, mpi::mpi_datatype< ScalarType >(), MPI_SUM, comm );
626
627 return dot_prod;
628}
629
630template < typename ScalarType, util::FlagLike FlagType >
631ScalarType masked_dot_product(
632 const grid::Grid4DDataScalar< ScalarType >& x,
633 const grid::Grid4DDataScalar< ScalarType >& y,
634 const grid::Grid4DDataScalar< FlagType >& mask,
635 const FlagType& mask_value,
636 MPI_Comm comm = MPI_COMM_WORLD )
637{
638 ScalarType dot_prod = 0.0;
639
640 Kokkos::parallel_reduce(
641 "masked_dot_product",
642 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
643 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_dot_prod ) {
644 const ScalarType mask_val = util::has_flag( mask( local_subdomain, i, j, k ), mask_value ) ? 1.0 : 0.0;
645 ScalarType val = x( local_subdomain, i, j, k ) * y( local_subdomain, i, j, k ) * mask_val;
646 local_dot_prod = local_dot_prod + val;
647 },
648 Kokkos::Sum< ScalarType >( dot_prod ) );
649
650 Kokkos::fence( "masked_dot_product" );
651
652 MPI_Allreduce( MPI_IN_PLACE, &dot_prod, 1, mpi::mpi_datatype< ScalarType >(), MPI_SUM, comm );
653
654 return dot_prod;
655}
656
657template < typename ScalarType, util::FlagLike FlagType, int VecDim >
658ScalarType masked_dot_product(
659 const grid::Grid4DDataVec< ScalarType, VecDim >& x,
660 const grid::Grid4DDataVec< ScalarType, VecDim >& y,
661 const grid::Grid4DDataScalar< FlagType >& mask,
662 const FlagType& mask_value,
663 MPI_Comm comm = MPI_COMM_WORLD )
664{
665 ScalarType dot_prod = 0.0;
666 for ( int d = 0; d < VecDim; ++d )
667 {
668 // masked_dot_product for scalar internally does MPI_Allreduce,
669 // so we accumulate the per-component results.
670 dot_prod += masked_dot_product( x.comp_[d], y.comp_[d], mask, mask_value, comm );
671 }
672 return dot_prod;
673}
674
675template < typename ScalarType >
676ScalarType dot_product_subset(
677 const grid::Grid4DDataScalar< ScalarType >& x,
678 const grid::Grid4DDataScalar< ScalarType >& y,
679 dense::Vec< int, 4 > start,
680 dense::Vec< int, 4 > end_excl,
681 MPI_Comm comm = MPI_COMM_WORLD )
682{
683 ScalarType dot_prod = 0.0;
684
685 Kokkos::parallel_reduce(
686 "masked_dot_product",
687 Kokkos::MDRangePolicy(
688 { start( 0 ), start( 1 ), start( 2 ), start( 3 ) },
689 { end_excl( 0 ), end_excl( 1 ), end_excl( 2 ), end_excl( 3 ) } ),
690 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, ScalarType& local_dot_prod ) {
691 ScalarType val = x( local_subdomain, i, j, k ) * y( local_subdomain, i, j, k );
692 local_dot_prod = local_dot_prod + val;
693 },
694 Kokkos::Sum< ScalarType >( dot_prod ) );
695
696 Kokkos::fence( "masked_dot_product" );
697
698 MPI_Allreduce( MPI_IN_PLACE, &dot_prod, 1, mpi::mpi_datatype< ScalarType >(), MPI_SUM, comm );
699
700 return dot_prod;
701}
702
703template < typename ScalarType >
704bool has_nan_or_inf( const grid::Grid4DDataScalar< ScalarType >& x, MPI_Comm comm = MPI_COMM_WORLD )
705{
706 bool has_nan_or_inf = false;
707
708 Kokkos::parallel_reduce(
709 "masked_dot_product",
710 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { x.extent( 0 ), x.extent( 1 ), x.extent( 2 ), x.extent( 3 ) } ),
711 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k, bool& local_has_nan_or_inf ) {
712 local_has_nan_or_inf = local_has_nan_or_inf || ( Kokkos::isnan( x( local_subdomain, i, j, k ) ) ||
713 Kokkos::isinf( x( local_subdomain, i, j, k ) ) );
714 },
715 Kokkos::LOr< bool >( has_nan_or_inf ) );
716
717 Kokkos::fence();
718
719 MPI_Allreduce( MPI_IN_PLACE, &has_nan_or_inf, 1, mpi::mpi_datatype< bool >(), MPI_LOR, comm );
720
721 return has_nan_or_inf;
722}
723
724template < typename ScalarType, int VecDim >
725bool has_nan_or_inf( const grid::Grid4DDataVec< ScalarType, VecDim >& x, MPI_Comm comm = MPI_COMM_WORLD )
726{
727 for ( int d = 0; d < VecDim; ++d )
728 {
729 if ( has_nan_or_inf( x.comp_[d], comm ) )
730 return true;
731 }
732 return false;
733}
734
735template < typename ScalarTypeDst, typename ScalarTypeSrc >
736void cast( const grid::Grid4DDataScalar< ScalarTypeDst >& dst, const grid::Grid4DDataScalar< ScalarTypeSrc >& src )
737{
738 Kokkos::parallel_for(
739 "cast",
740 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { dst.extent( 0 ), dst.extent( 1 ), dst.extent( 2 ), dst.extent( 3 ) } ),
741 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
742 dst( local_subdomain, i, j, k ) = static_cast< ScalarTypeDst >( src( local_subdomain, i, j, k ) );
743 } );
744
745 Kokkos::fence();
746}
747
748template < typename ScalarTypeDst >
749void rand( const grid::Grid4DDataScalar< ScalarTypeDst >& dst )
750{
751 static_assert(
752 std::is_same_v< ScalarTypeDst, double > || std::is_same_v< ScalarTypeDst, float >,
753 "Random integers not implemented. But can be done easily below." );
754
755 Kokkos::Random_XorShift64_Pool<> random_pool( /*seed=*/12345 );
756 Kokkos::parallel_for(
757 "rand",
758 Kokkos::MDRangePolicy( { 0, 0, 0, 0 }, { dst.extent( 0 ), dst.extent( 1 ), dst.extent( 2 ), dst.extent( 3 ) } ),
759 KOKKOS_LAMBDA( int local_subdomain, int i, int j, int k ) {
760 auto generator = random_pool.get_state();
761 dst( local_subdomain, i, j, k ) = static_cast< ScalarTypeDst >( generator.drand() );
762 random_pool.free_state( generator );
763 } );
764
765 Kokkos::fence();
766}
767
768template < typename ScalarTypeDst, int VecDim >
769void rand( const grid::Grid4DDataVec< ScalarTypeDst, VecDim >& dst )
770{
771 for ( int d = 0; d < VecDim; ++d )
772 rand( dst.comp_[d] );
773}
774
775} // namespace terra::kernels::common
bit_masking.hpp
bit_masks.hpp
grid_types.hpp
kokkos_wrapper.hpp
terra::grid::Grid5DDataScalar
Kokkos::View< ScalarType *****, Layout > Grid5DDataScalar
Definition grid_types.hpp:30
terra::grid::Grid3DDataScalar
Kokkos::View< ScalarType ***, Layout > Grid3DDataScalar
Definition grid_types.hpp:24
terra::grid::Grid4DDataScalar
Kokkos::View< ScalarType ****, Layout > Grid4DDataScalar
Definition grid_types.hpp:27
terra::grid::Grid2DDataScalar
Kokkos::View< ScalarType **, Layout > Grid2DDataScalar
Definition grid_types.hpp:21
terra::kernels::common
Definition grid_operations.hpp:9
terra::kernels::common::count_masked
ScalarType count_masked(const grid::Grid4DDataScalar< FlagType > &mask, const FlagType &mask_value, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:526
terra::kernels::common::extract_vector_component
void extract_vector_component(grid::Grid4DDataScalar< ScalarType > &component_out, const grid::Grid4DDataVec< ScalarType, VecDim > &vectorial_data_in, const int component)
Definition grid_operations.hpp:454
terra::kernels::common::invert_inplace
void invert_inplace(const grid::Grid4DDataScalar< ScalarType > &y)
Definition grid_operations.hpp:232
terra::kernels::common::max_abs_entry
ScalarType max_abs_entry(const grid::Grid4DDataScalar< ScalarType > &x, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:316
terra::kernels::common::has_nan_or_inf
bool has_nan_or_inf(const grid::Grid4DDataScalar< ScalarType > &x, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:704
terra::kernels::common::dot_product
ScalarType dot_product(const grid::Grid4DDataScalar< ScalarType > &x, const grid::Grid4DDataScalar< ScalarType > &y, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:608
terra::kernels::common::masked_dot_product
ScalarType masked_dot_product(const grid::Grid4DDataScalar< ScalarType > &x, const grid::Grid4DDataScalar< ScalarType > &y, const grid::Grid4DDataScalar< FlagType > &mask, const FlagType &mask_value, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:631
terra::kernels::common::max_abs_entry_subset
ScalarType max_abs_entry_subset(const grid::Grid4DDataScalar< ScalarType > &x, dense::Vec< int, 4 > start, dense::Vec< int, 4 > end_excl, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:377
terra::kernels::common::cast
void cast(const grid::Grid4DDataScalar< ScalarTypeDst > &dst, const grid::Grid4DDataScalar< ScalarTypeSrc > &src)
Definition grid_operations.hpp:736
terra::kernels::common::min_abs_entry
ScalarType min_abs_entry(const grid::Grid4DDataScalar< ScalarType > &x, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:296
terra::kernels::common::scale
void scale(const grid::Grid4DDataScalar< ScalarType > &x, ScalarType value)
Definition grid_operations.hpp:43
terra::kernels::common::vector_magnitude
void vector_magnitude(grid::Grid4DDataScalar< ScalarType > &magnitude_out, const grid::Grid4DDataVec< ScalarType, VecDim > &vectorial_data_in)
Definition grid_operations.hpp:428
terra::kernels::common::min_entry
ScalarType min_entry(const grid::Grid4DDataScalar< ScalarType > &x, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:276
terra::kernels::common::set_constant
void set_constant(const grid::Grid2DDataScalar< ScalarType > &x, ScalarType value)
Definition grid_operations.hpp:12
terra::kernels::common::mult_elementwise_inplace
void mult_elementwise_inplace(const grid::Grid4DDataScalar< ScalarType > &y, const grid::Grid4DDataScalar< ScalarType > &x)
Definition grid_operations.hpp:252
terra::kernels::common::max_vector_magnitude
ScalarType max_vector_magnitude(const grid::Grid4DDataVec< ScalarType, VecDim > &x, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:403
terra::kernels::common::sum_of_absolutes
ScalarType sum_of_absolutes(const grid::Grid4DDataScalar< ScalarType > &x, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:506
terra::kernels::common::assign_masked_else_keep_old
void assign_masked_else_keep_old(const grid::Grid4DDataScalar< ScalarType > &dst, const ScalarType &value, const grid::Grid4DDataScalar< FlagType > &mask_grid, const FlagType mask_value)
Definition grid_operations.hpp:54
terra::kernels::common::set_vector_component
void set_vector_component(grid::Grid4DDataVec< ScalarType, VecDim > &vectorial_data, const int component, const ScalarType constant)
Definition grid_operations.hpp:480
terra::kernels::common::rand
void rand(const grid::Grid4DDataScalar< ScalarTypeDst > &dst)
Definition grid_operations.hpp:749
terra::kernels::common::masked_sum
ScalarType masked_sum(const grid::Grid4DDataScalar< ScalarType > &x, const grid::Grid4DDataScalar< FlagType > &mask, const FlagType &mask_value, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:552
terra::kernels::common::lincomb
void lincomb(const grid::Grid4DDataScalar< ScalarType > &y, ScalarType c_0, ScalarType c_1, const grid::Grid4DDataScalar< ScalarType > &x_1)
Definition grid_operations.hpp:133
terra::kernels::common::dot_product_subset
ScalarType dot_product_subset(const grid::Grid4DDataScalar< ScalarType > &x, const grid::Grid4DDataScalar< ScalarType > &y, dense::Vec< int, 4 > start, dense::Vec< int, 4 > end_excl, MPI_Comm comm=MPI_COMM_WORLD)
Definition grid_operations.hpp:676
terra::mpi::mpi_datatype< bool >
MPI_Datatype mpi_datatype< bool >()
Definition mpi.hpp:231
terra::util::has_flag
constexpr bool has_flag(E mask_value, E flag) noexcept
Checks if a bitmask value contains a specific flag.
Definition bit_masking.hpp:43
terra::dense::Vec
Definition vec.hpp:9
terra::grid::Grid4DDataVec
SoA (Structure-of-Arrays) 4D vector grid data.
Definition grid_types.hpp:51
terra::grid::Grid4DDataVec::comp_
Grid4DDataScalar< ScalarType > comp_[VecDim]
Definition grid_types.hpp:57
terra::grid::Grid4DDataVec::extent
auto extent(int i) const
Definition grid_types.hpp:75