@@ -25,7 +25,6 @@ class ComputeE_FieldFunctor : public Base_Functor {
2525 Field &Ex, &Ey, &Ez;
2626 Field &Bx, &By, &Bz;
2727 Field &Jx, &Jy, &Jz;
28- int start_i, end_i;
2928 int Ni, Nj, Nk;
3029 FP current_coef;
3130 FP coef_dx, coef_dy, coef_dz;
@@ -36,12 +35,10 @@ class ComputeE_FieldFunctor : public Base_Functor {
3635 Field& Bx, Field& By, Field& Bz,
3736 Field& Jx, Field& Jy, Field& Jz,
3837 const FP & current_coef,
39- const int & start_i, const int & end_i,
4038 const int & Ni, const int & Nj, const int & Nk,
4139 const FP & coef_dx, const FP & coef_dy, const FP & coef_dz) :
4240 Ex (Ex), Ey(Ey), Ez(Ez), Bx(Bx), By(By), Bz(Bz),
4341 Jx (Jx), Jy(Jy), Jz(Jz), current_coef(current_coef),
44- start_i (start_i), end_i(end_i),
4542 Ni (Ni), Nj(Nj), Nk(Nk),
4643 coef_dx (coef_dx), coef_dy(coef_dy), coef_dz(coef_dz) {}
4744
@@ -52,125 +49,43 @@ class ComputeE_FieldFunctor : public Base_Functor {
5249 const FP & current_coef,
5350 const int bounds_i[2 ], const int bounds_j[2 ], const int bounds_k[2 ],
5451 const int & Ni, const int & Nj, const int & Nk,
55- const FP & coef_dx, const FP & coef_dy, FP & coef_dz) {
52+ const FP & coef_dx, const FP & coef_dy, const FP & coef_dz) {
5653
5754 ComputeE_FieldFunctor functor (
5855 Ex, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz,
59- current_coef, bounds_i[0 ], bounds_i[1 ],
60- Ni, Nj, Nk, coef_dx, coef_dy, coef_dz);
56+ current_coef, Ni, Nj, Nk, coef_dx, coef_dy, coef_dz);
6157
62- Kokkos::MDRangePolicy<Kokkos::Rank<2 >> policy (
63- {bounds_k[0 ], bounds_j[0 ]},
64- {bounds_k[1 ], bounds_j[1 ]});
58+ Kokkos::MDRangePolicy<Kokkos::Rank<3 , Kokkos::Iterate::Left, Kokkos::Iterate::Right >> policy (
59+ {bounds_k[0 ], bounds_j[0 ], bounds_i[ 0 ] },
60+ {bounds_k[1 ], bounds_j[1 ], bounds_i[ 1 ] });
6561
6662 Kokkos::parallel_for (" UpdateEField" , policy, functor);
6763 }
68- KOKKOS_INLINE_FUNCTION void operator ()(const int & k, const int & j) const {
64+ KOKKOS_INLINE_FUNCTION void operator ()(const int & k, const int & j, const int & i) const {
65+
66+ const int index = i + j * Ni + k * Ni * Nj;
6967
7068 int j_pred = j - 1 ;
7169 int k_pred = k - 1 ;
70+ int i_pred = i - 1 ;
7271
7372 applyPeriodicBoundary (k_pred, Nk);
7473 applyPeriodicBoundary (j_pred, Nj);
75-
76- const int index_kj_offset = j * Ni + k * Ni * Nj;
77- const int j_pred_kj_offset = j_pred * Ni + k * Ni * Nj;
78- const int k_pred_kj_offset = j * Ni + k_pred * Ni * Nj;
79-
80- int i_base = start_i;
81- for (; i_base + simd_width <= end_i; i_base += simd_width) {
82- const int current_simd_block_start_idx = i_base + index_kj_offset;
83- const int j_pred_simd_block_start_idx = i_base + j_pred_kj_offset;
84- const int k_pred_simd_block_start_idx = i_base + k_pred_kj_offset;
85-
86- simd_type Ex_simd, Ey_simd, Ez_simd;
87- simd_type Bx_simd, By_simd, Bz_simd;
88- simd_type Jx_simd, Jy_simd, Jz_simd;
89- simd_type Bz_pred_simd, Bx_j_pred_simd, By_pred_simd,
90- Bx_pred_simd, Bz_i_pred_simd, By_i_pred_simd;
91-
92- Ex_simd.copy_from (Ex.data () + current_simd_block_start_idx,
93- Kokkos::Experimental::simd_flag_default);
94- Ey_simd.copy_from (Ey.data () + current_simd_block_start_idx,
95- Kokkos::Experimental::simd_flag_default);
96- Ez_simd.copy_from (Ez.data () + current_simd_block_start_idx,
97- Kokkos::Experimental::simd_flag_default);
98-
99- Bx_simd.copy_from (Bx.data () + current_simd_block_start_idx,
100- Kokkos::Experimental::simd_flag_default);
101- By_simd.copy_from (By.data () + current_simd_block_start_idx,
102- Kokkos::Experimental::simd_flag_default);
103- Bz_simd.copy_from (Bz.data () + current_simd_block_start_idx,
104- Kokkos::Experimental::simd_flag_default);
105-
106- Jx_simd.copy_from (Jx.data () + current_simd_block_start_idx,
107- Kokkos::Experimental::simd_flag_default);
108- Jy_simd.copy_from (Jy.data () + current_simd_block_start_idx,
109- Kokkos::Experimental::simd_flag_default);
110- Jz_simd.copy_from (Jz.data () + current_simd_block_start_idx,
111- Kokkos::Experimental::simd_flag_default);
112-
113- Bz_pred_simd.copy_from (Bz.data () + j_pred_simd_block_start_idx,
114- Kokkos::Experimental::simd_flag_default); // Bz(i, j-1, k)
115- Bx_j_pred_simd.copy_from (Bx.data () + j_pred_simd_block_start_idx,
116- Kokkos::Experimental::simd_flag_default); // Bx(i, j-1, k)
117- By_pred_simd.copy_from (By.data () + k_pred_simd_block_start_idx,
118- Kokkos::Experimental::simd_flag_default); // By(i, j, k-1)
119- Bx_pred_simd.copy_from (Bx.data () + k_pred_simd_block_start_idx,
120- Kokkos::Experimental::simd_flag_default); // Bx(i, j, k-1)
121-
122- #pragma unroll
123- for (int lane = 0 ; lane < simd_width; ++lane) {
124- int i = i_base + lane;
125- int i_pred = i - 1 ;
126-
127- applyPeriodicBoundary (i_pred, Ni);
128-
129- int scalar_i_pred_idx = i_pred + index_kj_offset;
130-
131- Bz_i_pred_simd[lane] = Bz[scalar_i_pred_idx];
132- By_i_pred_simd[lane] = By[scalar_i_pred_idx];
133- }
134-
135- Ex_simd += current_coef * Jx_simd +
136- coef_dy * (Bz_simd - Bz_pred_simd) -
137- coef_dz * (By_simd - By_pred_simd);
138-
139- Ey_simd += current_coef * Jy_simd +
140- coef_dz * (Bx_simd - Bx_pred_simd) -
141- coef_dx * (Bz_simd - Bz_i_pred_simd);
142-
143- Ez_simd += current_coef * Jz_simd +
144- coef_dx * (By_simd - By_i_pred_simd) -
145- coef_dy * (Bx_simd - Bx_j_pred_simd);
146-
147- Ex_simd.copy_to (Ex.data () + current_simd_block_start_idx,
148- Kokkos::Experimental::simd_flag_default);
149- Ey_simd.copy_to (Ey.data () + current_simd_block_start_idx,
150- Kokkos::Experimental::simd_flag_default);
151- Ez_simd.copy_to (Ez.data () + current_simd_block_start_idx,
152- Kokkos::Experimental::simd_flag_default);
153- }
154- for (int i = i_base; i < end_i; ++i) {
155- const int index = i + index_kj_offset;
156- int i_pred = i - 1 ;
157-
158- applyPeriodicBoundary (i_pred, Ni);
159-
160- const int i_pred_idx = i_pred + index_kj_offset;
161- const int j_pred_idx = i + j_pred_kj_offset;
162- const int k_pred_idx = i + k_pred_kj_offset;
163-
164- Ex[index] += current_coef * Jx[index] +
165- coef_dy * (Bz[index] - Bz[j_pred_idx]) -
166- coef_dz * (By[index] - By[k_pred_idx]);
167- Ey[index] += current_coef * Jy[index] +
168- coef_dz * (Bx[index] - Bx[k_pred_idx]) -
169- coef_dx * (Bz[index] - Bz[i_pred_idx]);
170- Ez[index] += current_coef * Jz[index] +
171- coef_dx * (By[index] - By[i_pred_idx]) -
172- coef_dy * (Bx[index] - Bx[j_pred_idx]);
173- }
74+ applyPeriodicBoundary (i_pred, Ni);
75+
76+ const int i_pred_idx = i_pred + j * Ni + k * Ni * Nj;
77+ const int j_pred_idx = i + j_pred * Ni + k * Ni * Nj;
78+ const int k_pred_idx = i + j * Ni + k_pred * Ni * Nj;
79+
80+ Ex[index] += current_coef * Jx[index] +
81+ coef_dy * (Bz[index] - Bz[j_pred_idx]) -
82+ coef_dz * (By[index] - By[k_pred_idx]);
83+ Ey[index] += current_coef * Jy[index] +
84+ coef_dz * (Bx[index] - Bx[k_pred_idx]) -
85+ coef_dx * (Bz[index] - Bz[i_pred_idx]);
86+ Ez[index] += current_coef * Jz[index] +
87+ coef_dx * (By[index] - By[i_pred_idx]) -
88+ coef_dy * (Bx[index] - Bx[j_pred_idx]);
17489}
17590};
17691
@@ -180,16 +95,14 @@ class ComputeB_FieldFunctor : public Base_Functor {
18095 Field &Bx, &By, &Bz;
18196 int Ni, Nj, Nk;
18297 FP coef_dx, coef_dy, coef_dz;
183- int start_i, end_i;
18498public:
18599 ComputeB_FieldFunctor (
186100 Field& Ex, Field& Ey, Field& Ez,
187101 Field& Bx, Field& By, Field& Bz,
188- const int & start_i, const int & end_i,
189102 const int & Ni, const int & Nj, const int & Nk,
190103 const FP & coef_dx, const FP & coef_dy, const FP & coef_dz) :
191104 Ex (Ex), Ey(Ey), Ez(Ez), Bx(Bx), By(By), Bz(Bz),
192- Ni (Ni), Nj(Nj), Nk(Nk), start_i(start_i), end_i(end_i),
105+ Ni (Ni), Nj(Nj), Nk(Nk),
193106 coef_dx (coef_dx), coef_dy(coef_dy), coef_dz(coef_dz) {}
194107
195108 static void apply (
@@ -201,110 +114,37 @@ class ComputeB_FieldFunctor : public Base_Functor {
201114
202115 ComputeB_FieldFunctor functor (
203116 Ex, Ey, Ez, Bx, By, Bz,
204- bounds_i[0 ], bounds_i[1 ],
205117 Ni, Nj, Nk, coef_dx, coef_dy, coef_dz);
206118
207- Kokkos::MDRangePolicy<Kokkos::Rank<2 >> policy (
208- {bounds_k[0 ], bounds_j[0 ]},
209- {bounds_k[1 ], bounds_j[1 ]});
119+ Kokkos::MDRangePolicy<Kokkos::Rank<3 , Kokkos::Iterate::Left, Kokkos::Iterate::Right >> policy (
120+ {bounds_k[0 ], bounds_j[0 ], bounds_i[ 0 ] },
121+ {bounds_k[1 ], bounds_j[1 ], bounds_i[ 1 ] });
210122
211123 Kokkos::parallel_for (" UpdateBField" , policy, functor);
212124 }
213- KOKKOS_INLINE_FUNCTION void operator ()(const int & k, const int & j) const {
125+ KOKKOS_INLINE_FUNCTION void operator ()(const int & k, const int & j, const int & i) const {
126+ const int index = i + j * Ni + k * Ni * Nj;
127+
214128 int j_next = j + 1 ;
215129 int k_next = k + 1 ;
130+ int i_next = i + 1 ;
216131
217132 applyPeriodicBoundary (k_next, Nk);
218133 applyPeriodicBoundary (j_next, Nj);
134+ applyPeriodicBoundary (i_next, Ni);
219135
220- const int index_kj_offset = j * Ni + k * Ni * Nj;
221- const int j_next_kj_offset = j_next * Ni + k * Ni * Nj;
222- const int k_next_kj_offset = j * Ni + k_next * Ni * Nj;
223-
224- int i_base = start_i;
225- for (; i_base + simd_width <= end_i; i_base += simd_width) {
226- const int current_simd_block_start_idx = i_base + index_kj_offset;
227- const int j_next_simd_block_start_idx = i_base + j_next_kj_offset;
228- const int k_next_simd_block_start_idx = i_base + k_next_kj_offset;
229-
230- simd_type Bx_simd, By_simd, Bz_simd;
231- simd_type Ex_simd, Ey_simd, Ez_simd;
232- simd_type Ey_k_next_simd, Ex_j_next_simd;
233- simd_type Ez_i_next_simd, Ey_i_next_simd;
234- simd_type Ez_j_next_simd, Ex_k_next_simd;
235-
236- Bx_simd.copy_from (Bx.data () + current_simd_block_start_idx,
237- Kokkos::Experimental::simd_flag_default);
238- By_simd.copy_from (By.data () + current_simd_block_start_idx,
239- Kokkos::Experimental::simd_flag_default);
240- Bz_simd.copy_from (Bz.data () + current_simd_block_start_idx,
241- Kokkos::Experimental::simd_flag_default);
242-
243- Ex_simd.copy_from (Ex.data () + current_simd_block_start_idx,
244- Kokkos::Experimental::simd_flag_default);
245- Ey_simd.copy_from (Ey.data () + current_simd_block_start_idx,
246- Kokkos::Experimental::simd_flag_default);
247- Ez_simd.copy_from (Ez.data () + current_simd_block_start_idx,
248- Kokkos::Experimental::simd_flag_default);
249-
250- Ey_k_next_simd.copy_from (Ey.data () + k_next_simd_block_start_idx,
251- Kokkos::Experimental::simd_flag_default);
252- Ex_k_next_simd.copy_from (Ex.data () + k_next_simd_block_start_idx,
253- Kokkos::Experimental::simd_flag_default);
254-
255- Ex_j_next_simd.copy_from (Ex.data () + j_next_simd_block_start_idx,
256- Kokkos::Experimental::simd_flag_default);
257- Ez_j_next_simd.copy_from (Ez.data () + j_next_simd_block_start_idx,
258- Kokkos::Experimental::simd_flag_default);
259-
260- #pragma unroll
261- for (int lane = 0 ; lane < simd_width; ++lane) {
262- int i = i_base + lane;
263- int i_next = i + 1 ;
264- applyPeriodicBoundary (i_next, Ni);
265-
266- int scalar_i_next_idx = i_next + index_kj_offset;
267-
268- Ez_i_next_simd[lane] = Ez[scalar_i_next_idx];
269- Ey_i_next_simd[lane] = Ey[scalar_i_next_idx];
270- }
271-
272- Bx_simd += coef_dz * (Ey_k_next_simd - Ey_simd) -
273- coef_dy * (Ez_j_next_simd - Ez_simd);
274-
275- By_simd += coef_dx * (Ez_i_next_simd - Ez_simd) -
276- coef_dz * (Ex_k_next_simd - Ex_simd);
277-
278- Bz_simd += coef_dy * (Ex_j_next_simd - Ex_simd) -
279- coef_dx * (Ey_i_next_simd - Ey_simd);
280-
281- Bx_simd.copy_to (Bx.data () + current_simd_block_start_idx,
282- Kokkos::Experimental::simd_flag_default);
283- By_simd.copy_to (By.data () + current_simd_block_start_idx,
284- Kokkos::Experimental::simd_flag_default);
285- Bz_simd.copy_to (Bz.data () + current_simd_block_start_idx,
286- Kokkos::Experimental::simd_flag_default);
287- }
288-
289- for (int i = i_base; i < end_i; ++i) {
290- const int index = i + index_kj_offset;
136+ const int scalar_i_next_idx = i_next + j * Ni + k * Ni * Nj;
137+ const int scalar_j_next_idx = i + j_next * Ni + k * Ni * Nj;
138+ const int scalar_k_next_idx = i + j * Ni + k_next * Ni * Nj;
291139
292- int i_next = i + 1 ;
293- applyPeriodicBoundary (i_next, Ni);
294-
295- const int scalar_i_next_idx = i_next + index_kj_offset;
296- const int scalar_j_next_idx = i + j_next_kj_offset;
297- const int scalar_k_next_idx = i + k_next_kj_offset;
140+ Bx[index] += coef_dz * (Ey[scalar_k_next_idx] - Ey[index]) -
141+ coef_dy * (Ez[scalar_j_next_idx] - Ez[index]);
298142
299- Bx [index] += coef_dz * (Ey[scalar_k_next_idx ] - Ey [index]) -
300- coef_dy * (Ez[scalar_j_next_idx ] - Ez [index]);
143+ By [index] += coef_dx * (Ez[scalar_i_next_idx ] - Ez [index]) -
144+ coef_dz * (Ex[scalar_k_next_idx ] - Ex [index]);
301145
302- By[index] += coef_dx * (Ez[scalar_i_next_idx] - Ez[index]) -
303- coef_dz * (Ex[scalar_k_next_idx] - Ex[index]);
304-
305- Bz[index] += coef_dy * (Ex[scalar_j_next_idx] - Ex[index]) -
306- coef_dx * (Ey[scalar_i_next_idx] - Ey[index]);
307- }
146+ Bz[index] += coef_dy * (Ex[scalar_j_next_idx] - Ex[index]) -
147+ coef_dx * (Ey[scalar_i_next_idx] - Ey[index]);
308148}
309149};
310150
0 commit comments