Skip to content

Commit 514d315

Browse files
committed
Revert Kokkos test helper changes
1 parent d7b066e commit 514d315

1 file changed

Lines changed: 42 additions & 202 deletions

File tree

include/FDTD_kokkos/kokkos_functors.h

Lines changed: 42 additions & 202 deletions
Original file line numberDiff line numberDiff line change
@@ -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;
18498
public:
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

Comments
 (0)