deal.II version 9.8.0
\(\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\}}\)
Loading...
Searching...
No Matches
portable_tensor_product_kernels.h
Go to the documentation of this file.
1// -----------------------------------------------------------------------------
2//
3// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception OR LGPL-2.1-or-later
4// Copyright (C) 2017 - 2026 by the deal.II authors
5//
6// This file is part of the deal.II library.
7//
8// Detailed license information governing the source code and contributions
9// can be found in LICENSE.md and CONTRIBUTING.md at the top level directory.
10//
11// -----------------------------------------------------------------------------
12
13
14#ifndef dealii__tensor_product_kernels_h
15#define dealii__tensor_product_kernels_h
16
17#include <deal.II/base/config.h>
18
21
22#include <Kokkos_Core.hpp>
23
24
26
27
28namespace Portable
29{
30 namespace internal
31 {
36 // TODO: for now only the general variant is implemented
43
44
45
50 template <bool add, typename ViewTypeIn, typename ViewTypeOut>
53 const Kokkos::TeamPolicy<
54 MemorySpace::Default::kokkos_space::execution_space>::member_type
55 &team_member,
56 ViewTypeOut dst,
57 const ViewTypeIn src,
58 const int N)
59 {
60 Assert(dst.size() >= static_cast<unsigned int>(N), ExcInternalError());
61 Assert(src.size() >= static_cast<unsigned int>(N), ExcInternalError());
62 Kokkos::parallel_for(Kokkos::TeamVectorRange(team_member, N),
63 [&](const int i) {
64 if constexpr (add)
65 Kokkos::atomic_add(&dst(i), src(i));
66 else
67 dst(i) = src(i);
68 });
69
70 team_member.team_barrier();
71 }
72
73
74
75#if DEAL_II_KOKKOS_VERSION_GTE(4, 0, 0)
79 template <int n_rows,
80 int n_columns,
81 int direction,
82 typename Number,
83 typename ShapeDataMemorySpace,
84 bool contract_over_rows,
85 bool add,
86 typename ViewTypeIn,
87 typename ViewTypeOut>
89 apply_1d(const Kokkos::TeamPolicy<
90 MemorySpace::Default::kokkos_space::execution_space>::member_type
91 &team_member,
92 const Kokkos::View<Number *, ShapeDataMemorySpace> shape_data,
93 const ViewTypeIn in,
94 ViewTypeOut out)
95 {
96 constexpr int Nk = (contract_over_rows ? n_rows : n_columns),
97 Nq = (contract_over_rows ? n_columns : n_rows);
98
99 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
100 Assert(in.size() >= Nk, ExcInternalError());
101 Assert(out.size() >= Nq, ExcInternalError());
102
103 Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, Nq),
104 [&](const int q) {
105 Number sum = 0;
106 for (int k = 0; k < Nk; ++k)
107 {
108 const int shape_idx =
109 (contract_over_rows ? q + k * Nq :
110 k + q * Nk);
111 sum += shape_data(shape_idx) * in(k);
112 }
113
114 if constexpr (add)
115 Kokkos::atomic_add(&out(q), sum);
116 else
117 out(q) = sum;
118 });
119
120 team_member.team_barrier();
121 }
122
123
124
128 template <int n_rows,
129 int n_columns,
130 int direction,
131 typename Number,
132 typename ShapeDataMemorySpace,
133 bool contract_over_rows,
134 bool add,
135 typename ViewTypeIn,
136 typename ViewTypeOut>
138 apply_2d(const Kokkos::TeamPolicy<
139 MemorySpace::Default::kokkos_space::execution_space>::member_type
140 &team_member,
141 const Kokkos::View<Number *, ShapeDataMemorySpace> shape_data,
142 const ViewTypeIn in,
143 ViewTypeOut out)
144 {
145 using TeamType = Kokkos::TeamPolicy<
146 MemorySpace::Default::kokkos_space::execution_space>::member_type;
147
148 // Sizes of the input and output vectors:
149 // -----------------------------------------------------------
150 // direction | contract_over_rows | !contract_over_rows
151 // -----------------------------------------------------------
152 // 0 | m x m -> n x m | n x m -> m x m
153 // -----------------------------------------------------------
154 // 1 | n x m -> n x n | n x n -> n x m
155 // -----------------------------------------------------------
156 //
157 // Directions of the cycle indices:
158 // -----------------------------
159 // direction | j | q/k
160 // -----------------------------
161 // 0 | 1 | 0
162 // -----------------------------
163 // 1 | 0 | 1
164 // -----------------------------
165 constexpr int Nj = (direction < 1 ? n_rows : n_columns),
166 Nk = (contract_over_rows ? n_rows : n_columns),
167 Nq = (contract_over_rows ? n_columns : n_rows);
168
169 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
170 Assert(in.size() >= Nj * Nk, ExcInternalError());
171 Assert(out.size() >= Nj * Nq, ExcInternalError());
172
173 auto thread_policy =
174 Kokkos::TeamThreadMDRange<Kokkos::Rank<2>, TeamType>(team_member,
175 Nj,
176 Nq);
177 Kokkos::parallel_for(thread_policy, [&](const int j, const int q) {
178 const int base_shape = contract_over_rows ? q : q * n_columns;
179 const int stride_shape = contract_over_rows ? n_columns : 1;
180
181 const int base_in = (direction == 0 ? j * Nk : j);
182 const int stride_in = Utilities::pow(n_columns, direction);
183
184 Number sum = shape_data(base_shape) * in(base_in);
185 for (int k = 1; k < Nk; ++k)
186 sum += shape_data(base_shape + k * stride_shape) *
187 in(base_in + k * stride_in);
188
189 const int index_out = (direction == 0 ? j * Nq + q : j + q * Nj);
190
191 if constexpr (add)
192 Kokkos::atomic_add(&out(index_out), sum);
193 else
194 out(index_out) = sum;
195 });
196
197 team_member.team_barrier();
198 }
199
200
201
205 template <int n_rows,
206 int n_columns,
207 int direction,
208 typename Number,
209 typename ShapeDataMemorySpace,
210 bool contract_over_rows,
211 bool add,
212 typename ViewTypeIn,
213 typename ViewTypeOut>
215 apply_3d(const Kokkos::TeamPolicy<
216 MemorySpace::Default::kokkos_space::execution_space>::member_type
217 &team_member,
218 const Kokkos::View<Number *, ShapeDataMemorySpace> shape_data,
219 const ViewTypeIn in,
220 ViewTypeOut out)
221 {
222 using TeamType = Kokkos::TeamPolicy<
223 MemorySpace::Default::kokkos_space::execution_space>::member_type;
224
225 // Sizes of the input and output vectors:
226 // ------------------------------------------------------------------
227 // direction | contract_over_rows | !contract_over_rows
228 // ------------------------------------------------------------------
229 // 0 | m x m x m -> n x m x m | n x m x m -> m x m x m
230 // ------------------------------------------------------------------
231 // 1 | n x m x m -> n x n x m | n x n x m -> n x m x m
232 // ------------------------------------------------------------------
233 // 2 | n x n x m -> n x n x n | n x n x n -> n x n x m
234 // ------------------------------------------------------------------
235 //
236 // Directions of the cycle indices:
237 // -------------------------------------
238 // direction | i | j | q/k
239 // -------------------------------------
240 // 0 | 2 | 1 | 0
241 // -------------------------------------
242 // 1 | 0 | 2 | 1
243 // -------------------------------------
244 // 2 | 1 | 0 | 2
245 // -------------------------------------
246 constexpr int Ni = (direction < 1 ? n_rows : n_columns),
247 Nj = (direction < 2 ? n_rows : n_columns),
248 Nk = (contract_over_rows ? n_rows : n_columns),
249 Nq = (contract_over_rows ? n_columns : n_rows);
250
251 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
252 Assert(in.size() >= Ni * Nj * Nk, ExcInternalError());
253 Assert(out.size() >= Ni * Nj * Nq, ExcInternalError());
254
255 auto thread_policy = Kokkos::TeamThreadMDRange<Kokkos::Rank<3>, TeamType>(
256 team_member, Ni, Nj, Nq);
257 Kokkos::parallel_for(
258 thread_policy, [&](const int i, const int j, const int q) {
259 const int base_shape = contract_over_rows ? q : q * n_columns;
260 const int stride_shape = contract_over_rows ? n_columns : 1;
261
262 const int base_in =
263 (direction == 0 ? (i * Nj + j) * Nk :
264 (direction == 1 ? i + j * Ni * Nk : i * Nj + j));
265 const int stride_in = Utilities::pow(n_columns, direction);
266
267 Number sum = shape_data(base_shape) * in(base_in);
268 for (int k = 1; k < Nk; ++k)
269 sum += shape_data(base_shape + k * stride_shape) *
270 in(base_in + k * stride_in);
271
272 const int index_out =
273 (direction == 0 ? (i * Nj + j) * Nq + q :
274 (direction == 1 ? i + (j * Nq + q) * Ni :
275 (i + q * Ni) * Nj + j));
276
277 if constexpr (add)
278 Kokkos::atomic_add(&out(index_out), sum);
279 else
280 out(index_out) = sum;
281 });
282
283 team_member.team_barrier();
284 }
285#endif
286
287
288
289 template <int dim,
290 int n_rows,
291 int n_columns,
292 typename Number,
293 typename ShapeDataMemorySpace,
294 int direction,
295 bool contract_over_rows,
296 bool add,
297 typename ViewTypeIn,
298 typename ViewTypeOut>
300 apply(const Kokkos::TeamPolicy<
301 MemorySpace::Default::kokkos_space::execution_space>::member_type
302 &team_member,
303 const Kokkos::View<Number *, ShapeDataMemorySpace> shape_data,
304 const ViewTypeIn in,
305 ViewTypeOut out)
306 {
307 // We have two implementations for this apply() function. The first
308 // requires Kokkos version 4.0 or later, uses the modern
309 // TeamThreadMDRange, and is simpler. The second option (below)
310 // performs the i,j,k loop manually.
311 //
312 // The second implementation turns out to be slightly faster,
313 // at least until Kokkos 5.2. For now, always use the second
314 // implementation. The @ref step_104 "step-104" throughput for vmult() is 33 instead of 24
315 // MDoFs/s for an AMD W7800 and 135 instead of 96 MDoFs/s for an NVIDIA
316 // RTX 6000.
317#if 0
318 if constexpr (dim == 1)
319 apply_1d<n_rows,
320 n_columns,
321 direction,
322 Number,
323 ShapeDataMemorySpace,
324 contract_over_rows,
325 add>(team_member, shape_data, in, out);
326 if constexpr (dim == 2)
327 apply_2d<n_rows,
328 n_columns,
329 direction,
330 Number,
331 ShapeDataMemorySpace,
332 contract_over_rows,
333 add>(team_member, shape_data, in, out);
334 if constexpr (dim == 3)
335 apply_3d<n_rows,
336 n_columns,
337 direction,
338 Number,
339 ShapeDataMemorySpace,
340 contract_over_rows,
341 add>(team_member, shape_data, in, out);
342#else
343 // I: [0, m^{dim - direction - 1})
344 // J: [0, n^direction)
345 constexpr int NI = Utilities::pow(n_rows, dim - direction - 1);
346 constexpr int NJ = Utilities::pow(n_columns, direction);
347
348 constexpr int Nk = contract_over_rows ? n_rows : n_columns;
349 constexpr int Nq = contract_over_rows ? n_columns : n_rows;
350
351 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
352 Assert(in.size() >= NI * NJ * Nk, ExcInternalError());
353 Assert(out.size() >= NI * NJ * Nq, ExcInternalError());
354
355 constexpr int N = NI * NJ * Nq;
356 constexpr int stride = Utilities::pow(n_columns, direction);
357
358 Kokkos::parallel_for(
359 Kokkos::TeamThreadRange(team_member, N), [&](const int index_out) {
360 // index_in = (I Nk + k) n^direction + J
361 // index_out = (I Nq + q) n^direction + J
362 const int q = (index_out / stride) % Nq;
363 const int I = (index_out / stride) / Nq;
364 const int J = index_out % stride;
365
366 const int base_shape = contract_over_rows ? q : q * n_columns;
367 const int stride_shape = contract_over_rows ? n_columns : 1;
368 const int base_in = I * Nk * stride + J;
369
370 Number sum = shape_data(base_shape) * in(base_in);
371 for (int k = 1; k < Nk; ++k)
372 {
373 const int index_in = (I * Nk + k) * stride + J;
374 sum += shape_data(base_shape + k * stride_shape) * in(index_in);
375 }
376
377 if constexpr (add)
378 Kokkos::atomic_add(&out(index_out), sum);
379 else
380 out(index_out) = sum;
381 });
382
383 team_member.team_barrier();
384#endif
385 }
386
387
388
392 template <EvaluatorVariant variant,
393 int dim,
394 int n_rows,
395 int n_columns,
396 typename Number,
397 typename ShapeDataMemorySpace =
401
402
403
408 template <int dim,
409 int n_rows,
410 int n_columns,
411 typename Number,
412 typename ShapeDataMemorySpace>
414 dim,
415 n_rows,
416 n_columns,
417 Number,
418 ShapeDataMemorySpace>
419 {
420 public:
421 using TeamHandle = Kokkos::TeamPolicy<
422 MemorySpace::Default::kokkos_space::execution_space>::member_type;
423
424 using SharedView = Kokkos::View<Number *,
425 MemorySpace::Default::kokkos_space::
426 execution_space::scratch_memory_space,
427 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
428
431 const TeamHandle &team_member,
432 Kokkos::View<Number *, ShapeDataMemorySpace> shape_values,
433 Kokkos::View<Number *, ShapeDataMemorySpace> shape_gradients,
434 Kokkos::View<Number *, ShapeDataMemorySpace> co_shape_gradients,
436
441 template <int direction,
442 bool dof_to_quad,
443 bool add,
444 bool in_place,
445 typename ViewTypeIn,
446 typename ViewTypeOut>
448 values(const ViewTypeIn in, ViewTypeOut out) const;
449
454 template <int direction,
455 bool dof_to_quad,
456 bool add,
457 bool in_place,
458 typename ViewTypeIn,
459 typename ViewTypeOut>
461 gradients(const ViewTypeIn in, ViewTypeOut out) const;
462
467 template <int direction,
468 bool dof_to_quad,
469 bool add,
470 bool in_place,
471 typename ViewTypeIn,
472 typename ViewTypeOut>
474 co_gradients(const ViewTypeIn in, ViewTypeOut out) const;
475
480
484 Kokkos::View<Number *, ShapeDataMemorySpace> shape_values;
485
489 Kokkos::View<Number *, ShapeDataMemorySpace> shape_gradients;
490
494 Kokkos::View<Number *, ShapeDataMemorySpace> co_shape_gradients;
495
500 };
501
502
503
504 template <int dim,
505 int n_rows,
506 int n_columns,
507 typename Number,
508 typename ShapeDataMemorySpace>
511 dim,
512 n_rows,
513 n_columns,
514 Number,
515 ShapeDataMemorySpace>::
516 EvaluatorTensorProduct(
517 const TeamHandle &team_member,
518 Kokkos::View<Number *, ShapeDataMemorySpace> shape_values,
519 Kokkos::View<Number *, ShapeDataMemorySpace> shape_gradients,
520 Kokkos::View<Number *, ShapeDataMemorySpace> co_shape_gradients,
526 , temp(temp)
527 {}
528
529
530
531 template <int dim,
532 int n_rows,
533 int n_columns,
534 typename Number,
535 typename ShapeDataMemorySpace>
536 template <int direction,
537 bool dof_to_quad,
538 bool add,
539 bool in_place,
540 typename ViewTypeIn,
541 typename ViewTypeOut>
544 dim,
545 n_rows,
546 n_columns,
547 Number,
548 ShapeDataMemorySpace>::values(const ViewTypeIn in,
549 ViewTypeOut out) const
550 {
551 if constexpr (in_place)
552 {
553 apply<dim,
554 n_rows,
555 n_columns,
556 Number,
557 ShapeDataMemorySpace,
558 direction,
559 dof_to_quad,
560 false>(team_member, shape_values, in, temp);
561
562 populate_view<add>(team_member, out, temp, out.extent(0));
563 }
564 else
565 apply<dim,
566 n_rows,
567 n_columns,
568 Number,
569 ShapeDataMemorySpace,
570 direction,
571 dof_to_quad,
572 add>(team_member, shape_values, in, out);
573 }
574
575
576
577 template <int dim,
578 int n_rows,
579 int n_columns,
580 typename Number,
581 typename ShapeDataMemorySpace>
582 template <int direction,
583 bool dof_to_quad,
584 bool add,
585 bool in_place,
586 typename ViewTypeIn,
587 typename ViewTypeOut>
590 dim,
591 n_rows,
592 n_columns,
593 Number,
594 ShapeDataMemorySpace>::gradients(const ViewTypeIn in,
595 ViewTypeOut out)
596 const
597 {
598 if constexpr (in_place)
599 {
600 apply<dim,
601 n_rows,
602 n_columns,
603 Number,
604 ShapeDataMemorySpace,
605 direction,
606 dof_to_quad,
607 false>(team_member, shape_gradients, in, temp);
608
609 populate_view<add>(team_member, out, temp, out.extent(0));
610 }
611 else
612 apply<dim,
613 n_rows,
614 n_columns,
615 Number,
616 ShapeDataMemorySpace,
617 direction,
618 dof_to_quad,
619 add>(team_member, shape_gradients, in, out);
620 }
621
622
623
624 template <int dim,
625 int n_rows,
626 int n_columns,
627 typename Number,
628 typename ShapeDataMemorySpace>
629 template <int direction,
630 bool dof_to_quad,
631 bool add,
632 bool in_place,
633 typename ViewTypeIn,
634 typename ViewTypeOut>
638 dim,
639 n_rows,
640 n_columns,
641 Number,
642 ShapeDataMemorySpace>::co_gradients(const ViewTypeIn in,
643 ViewTypeOut out) const
644 {
645 if constexpr (in_place)
646 {
647 apply<dim,
648 n_columns,
649 n_columns,
650 Number,
651 ShapeDataMemorySpace,
652 direction,
653 dof_to_quad,
654 false>(team_member, co_shape_gradients, in, temp);
655
656 populate_view<add>(team_member, out, temp, out.extent(0));
657 }
658 else
659 apply<dim,
660 n_columns,
661 n_columns,
662 Number,
663 ShapeDataMemorySpace,
664 direction,
665 dof_to_quad,
666 add>(team_member, co_shape_gradients, in, out);
667 }
668 } // namespace internal
669} // namespace Portable
670
672
673#endif
#define DEAL_II_NAMESPACE_OPEN
Definition config.h:38
#define DEAL_II_HOST_DEVICE
Definition config.h:170
#define DEAL_II_NAMESPACE_CLOSE
Definition config.h:39
#define Assert(cond, exc)
static ::ExceptionBase & ExcInternalError()
*  *  *  RotationFunction< dim, Number >::RotationFunction Number(dim)
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 *, ShapeDataMemorySpace > 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:966
::Kokkos::DefaultExecutionSpace::memory_space kokkos_space
Kokkos::View< Number *, MemorySpace::Default::kokkos_space:: execution_space::scratch_memory_space, Kokkos::MemoryTraits< Kokkos::Unmanaged > > SharedView
EvaluatorTensorProduct(const TeamHandle &team_member, Kokkos::View< Number *, ShapeDataMemorySpace > shape_values, Kokkos::View< Number *, ShapeDataMemorySpace > shape_gradients, Kokkos::View< Number *, ShapeDataMemorySpace > co_shape_gradients, SharedView temp)