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