Loading [MathJax]/extensions/TeX/AMSsymbols.js
 deal.II version GIT relicensing-3075-gc235bd4825 2025-04-15 08:10:00+00:00
\(\newcommand{\dealvcentcolon}{\mathrel{\mathop{:}}}\) \(\newcommand{\dealcoloneq}{\dealvcentcolon\mathrel{\mkern-1.2mu}=}\) \(\newcommand{\jump}[1]{\left[\!\left[ #1 \right]\!\right]}\) \(\newcommand{\average}[1]{\left\{\!\left\{ #1 \right\}\!\right\}}\)
All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Modules Pages Concepts
portable_tensor_product_kernels.h
Go to the documentation of this file.
1// ------------------------------------------------------------------------
2//
3// SPDX-License-Identifier: LGPL-2.1-or-later
4// Copyright (C) 2017 - 2024 by the deal.II authors
5//
6// This file is part of the deal.II library.
7//
8// Part of the source code is dual licensed under Apache-2.0 WITH
9// LLVM-exception OR LGPL-2.1-or-later. Detailed license information
10// governing the source code and code contributions can be found in
11// LICENSE.md and CONTRIBUTING.md at the top level directory of deal.II.
12//
13// ------------------------------------------------------------------------
14
15
16#ifndef dealii__tensor_product_kernels_h
17#define dealii__tensor_product_kernels_h
18
19#include <deal.II/base/config.h>
20
23
24#include <Kokkos_Core.hpp>
25
26
28
29
30namespace Portable
31{
32 namespace internal
33 {
38 // TODO: for now only the general variant is implemented
45
46
47
52 template <bool add, typename ViewTypeIn, typename ViewTypeOut>
55 const Kokkos::TeamPolicy<
56 MemorySpace::Default::kokkos_space::execution_space>::member_type
57 &team_member,
58 ViewTypeOut dst,
59 const ViewTypeIn src,
60 const int N)
61 {
62 Assert(dst.size() >= N, ExcInternalError());
63 Assert(src.size() >= N, ExcInternalError());
64 Kokkos::parallel_for(Kokkos::TeamVectorRange(team_member, N),
65 [&](const int i) {
66 if constexpr (add)
67 Kokkos::atomic_add(&dst(i), src(i));
68 else
69 dst(i) = src(i);
70 });
71
72 team_member.team_barrier();
73 }
74
75
76
77#if DEAL_II_KOKKOS_VERSION_GTE(4, 0, 0)
81 template <int n_rows,
82 int n_columns,
83 int direction,
84 typename Number,
85 bool contract_over_rows,
86 bool add,
87 typename ViewTypeIn,
88 typename ViewTypeOut>
90 apply_1d(const Kokkos::TeamPolicy<
91 MemorySpace::Default::kokkos_space::execution_space>::member_type
92 &team_member,
93 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
94 shape_data,
95 const ViewTypeIn in,
96 ViewTypeOut out)
97 {
98 constexpr int Nk = (contract_over_rows ? n_rows : n_columns),
99 Nq = (contract_over_rows ? n_columns : n_rows);
100
101 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
102 Assert(in.size() >= Nk, ExcInternalError());
103 Assert(out.size() >= Nq, ExcInternalError());
104
105 Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, Nq),
106 [&](const int q) {
107 Number sum = 0;
108 for (int k = 0; k < Nk; ++k)
109 {
110 const int shape_idx =
111 (contract_over_rows ? q + k * Nq :
112 k + q * Nk);
113 sum += shape_data(shape_idx) * in(k);
114 }
115
116 if constexpr (add)
117 Kokkos::atomic_add(&out(q), sum);
118 else
119 out(q) = sum;
120 });
121
122 team_member.team_barrier();
123 }
124
125
126
130 template <int n_rows,
131 int n_columns,
132 int direction,
133 typename Number,
134 bool contract_over_rows,
135 bool add,
136 typename ViewTypeIn,
137 typename ViewTypeOut>
139 apply_2d(const Kokkos::TeamPolicy<
140 MemorySpace::Default::kokkos_space::execution_space>::member_type
141 &team_member,
142 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
143 shape_data,
144 const ViewTypeIn in,
145 ViewTypeOut out)
146 {
147 using TeamType = Kokkos::TeamPolicy<
148 MemorySpace::Default::kokkos_space::execution_space>::member_type;
149
150 // Sizes of the input and output vectors:
151 // -----------------------------------------------------------
152 // direction | contract_over_rows | !contract_over_rows
153 // -----------------------------------------------------------
154 // 0 | m x m -> n x m | n x m -> m x m
155 // -----------------------------------------------------------
156 // 1 | n x m -> n x n | n x n -> n x m
157 // -----------------------------------------------------------
158 //
159 // Directions of the cycle indices:
160 // -----------------------------
161 // direction | j | q/k
162 // -----------------------------
163 // 0 | 1 | 0
164 // -----------------------------
165 // 1 | 0 | 1
166 // -----------------------------
167 constexpr int Nj = (direction < 1 ? n_rows : n_columns),
168 Nk = (contract_over_rows ? n_rows : n_columns),
169 Nq = (contract_over_rows ? n_columns : n_rows);
170
171 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
172 Assert(in.size() >= Nj * Nk, ExcInternalError());
173 Assert(out.size() >= Nj * Nq, ExcInternalError());
174
175 auto thread_policy =
176 Kokkos::TeamThreadMDRange<Kokkos::Rank<2>, TeamType>(team_member,
177 Nj,
178 Nq);
179 Kokkos::parallel_for(thread_policy, [&](const int j, const int q) {
180 const int base_shape = contract_over_rows ? q : q * n_columns;
181 const int stride_shape = contract_over_rows ? n_columns : 1;
182
183 const int base_in = (direction == 0 ? j * Nk : j);
184 const int stride_in = Utilities::pow(n_columns, direction);
185
186 Number sum = shape_data(base_shape) * in(base_in);
187 for (int k = 1; k < Nk; ++k)
188 sum += shape_data(base_shape + k * stride_shape) *
189 in(base_in + k * stride_in);
190
191 const int index_out = (direction == 0 ? j * Nq + q : j + q * Nj);
192
193 if constexpr (add)
194 Kokkos::atomic_add(&out(index_out), sum);
195 else
196 out(index_out) = sum;
197 });
198
199 team_member.team_barrier();
200 }
201
202
203
207 template <int n_rows,
208 int n_columns,
209 int direction,
210 typename Number,
211 bool contract_over_rows,
212 bool add,
213 typename ViewTypeIn,
214 typename ViewTypeOut>
216 apply_3d(const Kokkos::TeamPolicy<
217 MemorySpace::Default::kokkos_space::execution_space>::member_type
218 &team_member,
219 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
220 shape_data,
221 const ViewTypeIn in,
222 ViewTypeOut out)
223 {
224 using TeamType = Kokkos::TeamPolicy<
225 MemorySpace::Default::kokkos_space::execution_space>::member_type;
226
227 // Sizes of the input and output vectors:
228 // ------------------------------------------------------------------
229 // direction | contract_over_rows | !contract_over_rows
230 // ------------------------------------------------------------------
231 // 0 | m x m x m -> n x m x m | n x m x m -> m x m x m
232 // ------------------------------------------------------------------
233 // 1 | n x m x m -> n x n x m | n x n x m -> n x m x m
234 // ------------------------------------------------------------------
235 // 2 | n x n x m -> n x n x n | n x n x n -> n x n x m
236 // ------------------------------------------------------------------
237 //
238 // Directions of the cycle indices:
239 // -------------------------------------
240 // direction | i | j | q/k
241 // -------------------------------------
242 // 0 | 2 | 1 | 0
243 // -------------------------------------
244 // 1 | 0 | 2 | 1
245 // -------------------------------------
246 // 2 | 1 | 0 | 2
247 // -------------------------------------
248 constexpr int Ni = (direction < 1 ? n_rows : n_columns),
249 Nj = (direction < 2 ? n_rows : n_columns),
250 Nk = (contract_over_rows ? n_rows : n_columns),
251 Nq = (contract_over_rows ? n_columns : n_rows);
252
253 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
254 Assert(in.size() >= Ni * Nj * Nk, ExcInternalError());
255 Assert(out.size() >= Ni * Nj * Nq, ExcInternalError());
256
257 auto thread_policy = Kokkos::TeamThreadMDRange<Kokkos::Rank<3>, TeamType>(
258 team_member, Ni, Nj, Nq);
259 Kokkos::parallel_for(
260 thread_policy, [&](const int i, const int j, const int q) {
261 const int base_shape = contract_over_rows ? q : q * n_columns;
262 const int stride_shape = contract_over_rows ? n_columns : 1;
263
264 const int base_in =
265 (direction == 0 ? (i * Nj + j) * Nk :
266 (direction == 1 ? i + j * Ni * Nk : i * Nj + j));
267 const int stride_in = Utilities::pow(n_columns, direction);
268
269 Number sum = shape_data(base_shape) * in(base_in);
270 for (int k = 1; k < Nk; ++k)
271 sum += shape_data(base_shape + k * stride_shape) *
272 in(base_in + k * stride_in);
273
274 const int index_out =
275 (direction == 0 ? (i * Nj + j) * Nq + q :
276 (direction == 1 ? i + (j * Nq + q) * Ni :
277 (i + q * Ni) * Nj + j));
278
279 if constexpr (add)
280 Kokkos::atomic_add(&out(index_out), sum);
281 else
282 out(index_out) = sum;
283 });
284
285 team_member.team_barrier();
286 }
287#endif
288
289
290
291 template <int dim,
292 int n_rows,
293 int n_columns,
294 typename Number,
295 int direction,
296 bool contract_over_rows,
297 bool add,
298 typename ViewTypeIn,
299 typename ViewTypeOut>
301 apply(const Kokkos::TeamPolicy<
302 MemorySpace::Default::kokkos_space::execution_space>::member_type
303 &team_member,
304 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
305 shape_data,
306 const ViewTypeIn in,
307 ViewTypeOut out)
308 {
309#if DEAL_II_KOKKOS_VERSION_GTE(4, 0, 0)
310 if constexpr (dim == 1)
311 apply_1d<n_rows, n_columns, direction, Number, contract_over_rows, add>(
312 team_member, shape_data, in, out);
313 if constexpr (dim == 2)
314 apply_2d<n_rows, n_columns, direction, Number, contract_over_rows, add>(
315 team_member, shape_data, in, out);
316 if constexpr (dim == 3)
317 apply_3d<n_rows, n_columns, direction, Number, contract_over_rows, add>(
318 team_member, shape_data, in, out);
319#else
320 // I: [0, m^{dim - direction - 1})
321 // J: [0, n^direction)
322 constexpr int NI = Utilities::pow(n_rows, dim - direction - 1);
323 constexpr int NJ = Utilities::pow(n_columns, direction);
324
325 constexpr int Nk = contract_over_rows ? n_rows : n_columns;
326 constexpr int Nq = contract_over_rows ? n_columns : n_rows;
327
328 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
329 Assert(in.size() >= NI * NJ * Nk, ExcInternalError());
330 Assert(out.size() >= NI * NJ * Nq, ExcInternalError());
331
332 constexpr int N = NI * NJ * Nq;
333 constexpr int stride = Utilities::pow(n_columns, direction);
334
335 Kokkos::parallel_for(
336 Kokkos::TeamThreadRange(team_member, N), [&](const int index_out) {
337 // index_in = (I Nk + k) n^direction + J
338 // index_out = (I Nq + q) n^direction + J
339 const int q = (index_out / stride) % Nq;
340 const int I = (index_out / stride) / Nq;
341 const int J = index_out % stride;
342
343 const int base_shape = contract_over_rows ? q : q * n_columns;
344 const int stride_shape = contract_over_rows ? n_columns : 1;
345 const int base_in = I * Nk * stride + J;
346
347 Number sum = shape_data(base_shape) * in(base_in);
348 for (int k = 1; k < Nk; ++k)
349 {
350 const int index_in = (I * Nk + k) * stride + J;
351 sum += shape_data(base_shape + k * stride_shape) * in(index_in);
352 }
353
354 if constexpr (add)
355 Kokkos::atomic_add(&out(index_out), sum);
356 else
357 out(index_out) = sum;
358 });
359
360 team_member.team_barrier();
361#endif
362 }
363
364
365
369 template <EvaluatorVariant variant,
370 int dim,
371 int n_rows,
372 int n_columns,
373 typename Number>
376
377
378
383 template <int dim, int n_rows, int n_columns, typename Number>
385 dim,
386 n_rows,
387 n_columns,
388 Number>
389 {
390 public:
391 using TeamHandle = Kokkos::TeamPolicy<
392 MemorySpace::Default::kokkos_space::execution_space>::member_type;
393
394 using SharedView = Kokkos::View<Number *,
395 MemorySpace::Default::kokkos_space::
396 execution_space::scratch_memory_space,
397 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
398
401 const TeamHandle &team_member,
402 Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values,
403 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
404 shape_gradients,
405 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
406 co_shape_gradients,
407 SharedView temp);
408
413 template <int direction,
414 bool dof_to_quad,
415 bool add,
416 bool in_place,
417 typename ViewTypeIn,
418 typename ViewTypeOut>
420 values(const ViewTypeIn in, ViewTypeOut out) const;
421
426 template <int direction,
427 bool dof_to_quad,
428 bool add,
429 bool in_place,
430 typename ViewTypeIn,
431 typename ViewTypeOut>
433 gradients(const ViewTypeIn in, ViewTypeOut out) const;
434
439 template <int direction,
440 bool dof_to_quad,
441 bool add,
442 bool in_place,
443 typename ViewTypeIn,
444 typename ViewTypeOut>
446 co_gradients(const ViewTypeIn in, ViewTypeOut out) const;
447
452
456 Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values;
457
461 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
463
467 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
469
474 };
475
476
477
478 template <int dim, int n_rows, int n_columns, typename Number>
482 const TeamHandle &team_member,
483 Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values,
484 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
485 shape_gradients,
486 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
487 co_shape_gradients,
488 SharedView temp)
489 : team_member(team_member)
490 , shape_values(shape_values)
491 , shape_gradients(shape_gradients)
492 , co_shape_gradients(co_shape_gradients)
493 , temp(temp)
494 {}
495
496
497
498 template <int dim, int n_rows, int n_columns, typename Number>
499 template <int direction,
500 bool dof_to_quad,
501 bool add,
502 bool in_place,
503 typename ViewTypeIn,
504 typename ViewTypeOut>
507 values(const ViewTypeIn in, ViewTypeOut out) const
508 {
509 if constexpr (in_place)
510 {
511 apply<dim, n_rows, n_columns, Number, direction, dof_to_quad, false>(
512 team_member, shape_values, in, temp);
513
514 populate_view<add>(team_member, out, temp, out.extent(0));
515 }
516 else
517 apply<dim, n_rows, n_columns, Number, direction, dof_to_quad, add>(
518 team_member, shape_values, in, out);
519 }
520
521
522
523 template <int dim, int n_rows, int n_columns, typename Number>
524 template <int direction,
525 bool dof_to_quad,
526 bool add,
527 bool in_place,
528 typename ViewTypeIn,
529 typename ViewTypeOut>
532 gradients(const ViewTypeIn in, ViewTypeOut out) const
533 {
534 if constexpr (in_place)
535 {
536 apply<dim, n_rows, n_columns, Number, direction, dof_to_quad, false>(
537 team_member, shape_gradients, in, temp);
538
539 populate_view<add>(team_member, out, temp, out.extent(0));
540 }
541 else
542 apply<dim, n_rows, n_columns, Number, direction, dof_to_quad, add>(
543 team_member, shape_gradients, in, out);
544 }
545
546
547
548 template <int dim, int n_rows, int n_columns, typename Number>
549 template <int direction,
550 bool dof_to_quad,
551 bool add,
552 bool in_place,
553 typename ViewTypeIn,
554 typename ViewTypeOut>
557 co_gradients(const ViewTypeIn in, ViewTypeOut out) const
558 {
559 if constexpr (in_place)
560 {
561 apply<dim,
562 n_columns,
563 n_columns,
564 Number,
565 direction,
566 dof_to_quad,
567 false>(team_member, co_shape_gradients, in, temp);
568
569 populate_view<add>(team_member, out, temp, out.extent(0));
570 }
571 else
572 apply<dim, n_columns, n_columns, Number, direction, dof_to_quad, add>(
573 team_member, co_shape_gradients, in, out);
574 }
575 } // namespace internal
576} // namespace Portable
577
579
580#endif
#define DEAL_II_NAMESPACE_OPEN
Definition config.h:35
#define DEAL_II_HOST_DEVICE
Definition config.h:166
#define DEAL_II_NAMESPACE_CLOSE
Definition config.h:36
#define Assert(cond, exc)
static ::ExceptionBase & ExcInternalError()
void populate_view(const Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type &team_member, ViewTypeOut dst, const ViewTypeIn src, const int N)
void apply(const Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type &team_member, const Kokkos::View< Number *, MemorySpace::Default::kokkos_space > shape_data, const ViewTypeIn in, ViewTypeOut out)
T sum(const T &t, const MPI_Comm mpi_communicator)
constexpr T pow(const T base, const int iexp)
Definition utilities.h:967
Kokkos::View< Number *, MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space, Kokkos::MemoryTraits< Kokkos::Unmanaged > > SharedView
Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type TeamHandle