deal.II version GIT relicensing-6809-ge913b9bb34 2026-09-25 17:20:01+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\}}\)
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
19#include <deal.II/base/memory_space.h>
20#include <deal.II/base/utilities.h>
21
22#include <Kokkos_Core.hpp>
23
24
25DEAL_II_NAMESPACE_OPEN
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>
51 DEAL_II_HOST_DEVICE void
52 populate_view(
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>
88 DEAL_II_HOST_DEVICE void
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>
137 DEAL_II_HOST_DEVICE void
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>
214 DEAL_II_HOST_DEVICE void
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>
299 DEAL_II_HOST_DEVICE void
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 =
398 MemorySpace::Default::kokkos_space>
401
402
403
408 template <int dim,
409 int n_rows,
410 int n_columns,
411 typename Number,
412 typename ShapeDataMemorySpace>
413 struct EvaluatorTensorProduct<evaluate_general,
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
429 DEAL_II_HOST_DEVICE
430 EvaluatorTensorProduct(
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,
435 SharedView temp);
436
441 template <int direction,
442 bool dof_to_quad,
443 bool add,
444 bool in_place,
445 typename ViewTypeIn,
446 typename ViewTypeOut>
447 DEAL_II_HOST_DEVICE void
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>
460 DEAL_II_HOST_DEVICE void
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>
473 DEAL_II_HOST_DEVICE void
474 co_gradients(const ViewTypeIn in, ViewTypeOut out) const;
475
479 const TeamHandle &team_member;
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
499 SharedView temp;
500 };
501
502
503
504 template <int dim,
505 int n_rows,
506 int n_columns,
507 typename Number,
508 typename ShapeDataMemorySpace>
509 DEAL_II_HOST_DEVICE
510 EvaluatorTensorProduct<evaluate_general,
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,
521 SharedView temp)
522 : team_member(team_member)
523 , shape_values(shape_values)
524 , shape_gradients(shape_gradients)
525 , co_shape_gradients(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>
542 DEAL_II_HOST_DEVICE void
543 EvaluatorTensorProduct<evaluate_general,
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>
588 DEAL_II_HOST_DEVICE void
589 EvaluatorTensorProduct<evaluate_general,
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>
635 DEAL_II_HOST_DEVICE void
636 EvaluatorTensorProduct<
637 evaluate_general,
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
669 namespace batched
670 {
683 template <int n_rows,
684 int n_columns,
685 bool contract_over_rows,
686 bool add,
687 int stride_in,
688 int stride_out,
689 typename TypeMatrix,
690 typename TypeIn,
691 typename TypeOut>
692 DEAL_II_HOST_DEVICE inline void
693 apply_matrix_vector_product(const TypeMatrix matrix,
694 const TypeIn in,
695 TypeOut out)
696 {
697 constexpr int mm = contract_over_rows ? n_rows : n_columns;
698 constexpr int nn = contract_over_rows ? n_columns : n_rows;
699
700 // Deduce the type of the underlying TypeOut array type. This is
701 // needed for caching below
702 using Number =
703 std::remove_cv_t<std::remove_reference_t<decltype(out[0])>>;
704
705 // Cache the input in registers once, rather than re-reading
706 // shared memory for every output index below.
707 Number r_in[mm];
708 for (int k = 0; k < mm; ++k)
709 r_in[k] = in[k * stride_in];
710
711 for (int q = 0; q < nn; ++q)
712 {
713 Number sum = 0;
714 for (int k = 0; k < mm; ++k)
715 {
716 const int row = contract_over_rows ? k : q;
717 const int col = contract_over_rows ? q : k;
718 sum += matrix[row * n_columns + col] * r_in[k];
719 }
720
721 if constexpr (add)
722 out[q * stride_out] += sum;
723 else
724 out[q * stride_out] = sum;
725 }
726 }
727
732 template <int n_rows,
733 int n_columns,
734 bool contract_over_rows,
735 bool add,
736 typename TypeMatrix,
737 typename TypeIn,
738 typename TypeOut>
739 DEAL_II_HOST_DEVICE inline void
740 apply_matrix_vector_product(const TypeMatrix matrix,
741 const TypeIn in,
742 TypeOut out,
743 const int stride_in,
744 const int stride_out)
745 {
746 constexpr int mm = contract_over_rows ? n_rows : n_columns;
747 constexpr int nn = contract_over_rows ? n_columns : n_rows;
748
749 // Deduce the type of the underlying ViewTypeOut array type. This is
750 // needed for caching below
751 using Number =
752 std::remove_cv_t<std::remove_reference_t<decltype(out[0])>>;
753
754 // Cache the input in registers once, rather than re-reading
755 // shared memory for every output index below.
756 Number r_in[mm];
757 for (int k = 0; k < mm; ++k)
758 r_in[k] = in[k * stride_in];
759
760 for (int q = 0; q < nn; ++q)
761 {
762 Number sum = 0;
763 for (int k = 0; k < mm; ++k)
764 {
765 const int row = contract_over_rows ? k : q;
766 const int col = contract_over_rows ? q : k;
767 sum += matrix[row * n_columns + col] * r_in[k];
768 }
769
770 if constexpr (add)
771 out[q * stride_out] += sum;
772 else
773 out[q * stride_out] = sum;
774 }
775 }
776
808 template <
809 int dim,
810 int direction,
811 int n_rows,
812 int n_columns,
813 bool contract_over_rows,
814 bool add,
815 typename ViewTypeMatrix,
816 typename ViewTypeIn,
817 typename ViewTypeOut,
818 typename = std::enable_if_t<Kokkos::is_view<ViewTypeOut>::value>>
819 DEAL_II_HOST_DEVICE inline void
820 apply(const Kokkos::TeamPolicy<
821 MemorySpace::Default::kokkos_space::execution_space>::member_type
822 &team_member,
823 const ViewTypeMatrix shape_data,
824 const ViewTypeIn in,
825 ViewTypeOut out,
826 const int batch_size = 1)
827 {
828 static_assert(direction >= 0 && direction < dim,
829 "direction must be in [0, dim)");
830 Assert(shape_data.size() == n_rows * n_columns, ExcInternalError());
831
832 constexpr int mm = contract_over_rows ? n_rows : n_columns;
833 constexpr int nn = contract_over_rows ? n_columns : n_rows;
834
835 // combined extent of the already-transformed axes
836 // (role < direction, extent n_columns each)
837 constexpr int n_blocks1 = Utilities::pow(n_columns, direction);
838
839 // n_blocks2: combined extent ofthe not-yet-transformed
840 // axes (role > direction, extent n_rows each)
841 constexpr int n_blocks2 = Utilities::pow(n_rows, dim - direction - 1);
842
843 constexpr int n_in_per_elmt = n_blocks1 * mm * n_blocks2;
844 constexpr int n_out_per_elmt = n_blocks1 * nn * n_blocks2;
845
846 // Unlike apply_1d()/apply_2d()/apply_3d() above (one cell per call),
847 // this is the batched variant -- batch_size cells'
848 // worth of data, laid out contiguously per element, share one in/out
849 // view, so the required size scales with the batch size too.
850 Assert(in.size() >=
851 static_cast<std::size_t>(batch_size) * n_in_per_elmt,
852 ExcInternalError());
853 Assert(out.size() >=
854 static_cast<std::size_t>(batch_size) * n_out_per_elmt,
855 ExcInternalError());
856
857 Kokkos::parallel_for(
858 Kokkos::TeamVectorRange(team_member,
859 batch_size * n_blocks1 * n_blocks2),
860 [&](const int tid) {
861 const int e = tid / (n_blocks1 * n_blocks2);
862 const int rem = tid % (n_blocks1 * n_blocks2);
863 const int i2 = rem / n_blocks1;
864 const int i1 = rem % n_blocks1;
865
866 const int in_offset = e * n_in_per_elmt + i2 * n_blocks1 * mm + i1;
867 const int out_offset =
868 e * n_out_per_elmt + i2 * n_blocks1 * nn + i1;
869
870 apply_matrix_vector_product<n_rows,
871 n_columns,
872 contract_over_rows,
873 add,
874 n_blocks1,
875 n_blocks1>(
876 shape_data,
877 Kokkos::subview(
878 in,
879 Kokkos::make_pair(in_offset, static_cast<int>(in.extent(0)))),
880 Kokkos::subview(out,
881 Kokkos::make_pair(
882 out_offset, static_cast<int>(out.extent(0)))));
883 });
884
885 team_member.team_barrier();
886 }
887
892 template <
893 bool add,
894 typename ViewTypeOut,
895 typename ViewTypeIn,
896 typename = std::enable_if_t<Kokkos::is_view<ViewTypeOut>::value>>
897 DEAL_II_HOST_DEVICE inline void
898 populate_view(
899 const Kokkos::TeamPolicy<
900 MemorySpace::Default::kokkos_space::execution_space>::member_type
901 &team_member,
902 ViewTypeOut dst,
903 const ViewTypeIn src,
904 const int N)
905 {
906 Kokkos::parallel_for(Kokkos::TeamVectorRange(team_member, N),
907 [&](const int tid) {
908 if constexpr (add)
909 dst(tid) += src(tid);
910 else
911 dst(tid) = src(tid);
912 });
913
914 team_member.team_barrier();
915 }
916
917
921 template <EvaluatorVariant variant,
922 int dim,
923 int n_rows,
924 int n_columns,
925 typename Number>
928
929
934 template <int dim, int n_rows, int n_columns, typename Number>
935 struct EvaluatorTensorProduct<evaluate_general,
936 dim,
937 n_rows,
938 n_columns,
939 Number>
940 {
941 public:
942 using TeamHandle = Kokkos::TeamPolicy<
943 MemorySpace::Default::kokkos_space::execution_space>::member_type;
944
945 using ShapeDataType =
946 Kokkos::View<Number *,
947 MemorySpace::Default::kokkos_space::execution_space::
948 scratch_memory_space,
949 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
950
951 using SharedView =
952 Kokkos::View<Number *,
953 MemorySpace::Default::kokkos_space::execution_space::
954 scratch_memory_space,
955 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
956
957
958 DEAL_II_HOST_DEVICE
959 EvaluatorTensorProduct(const TeamHandle &team_member,
960 ShapeDataType shape_values,
961 ShapeDataType shape_gradients,
962 ShapeDataType co_shape_gradients,
963 SharedView temp,
964 const int batch_size = 1);
965
970 template <int direction,
971 bool dof_to_quad,
972 bool add,
973 bool in_place,
974 typename ViewTypeIn,
975 typename ViewTypeOut>
976 DEAL_II_HOST_DEVICE void
977 value(const ViewTypeIn in, ViewTypeOut out) const;
978
983 template <int direction,
984 bool dof_to_quad,
985 bool add,
986 bool in_place,
987 typename ViewTypeIn,
988 typename ViewTypeOut>
989 DEAL_II_HOST_DEVICE void
990 derivative(const ViewTypeIn in, ViewTypeOut out) const;
991
996 template <int direction,
997 bool dof_to_quad,
998 bool add,
999 bool in_place,
1000 typename ViewTypeIn,
1001 typename ViewTypeOut>
1002 DEAL_II_HOST_DEVICE void
1003 co_derivative(const ViewTypeIn in, ViewTypeOut out) const;
1004
1005
1020 template <bool transpose,
1021 bool add,
1022 typename ViewTypeIn,
1023 typename ViewTypeOut>
1024 DEAL_II_HOST_DEVICE void
1025 co_gradient(const ViewTypeIn in, ViewTypeOut out) const;
1026
1027 private:
1028 const TeamHandle &team_member;
1029 ShapeDataType shape_values;
1030 ShapeDataType shape_gradients;
1031 ShapeDataType co_shape_gradients;
1032 SharedView temp;
1033 const int batch_size;
1034 };
1035
1036 template <int dim, int n_rows, int n_columns, typename Number>
1037 DEAL_II_HOST_DEVICE
1038 EvaluatorTensorProduct<evaluate_general, dim, n_rows, n_columns, Number>::
1039 EvaluatorTensorProduct(const TeamHandle &team_member,
1040 ShapeDataType shape_values,
1041 ShapeDataType shape_gradients,
1042 ShapeDataType co_shape_gradients,
1043 SharedView temp,
1044 const int batch_size)
1045 : team_member(team_member)
1046 , shape_values(shape_values)
1047 , shape_gradients(shape_gradients)
1048 , co_shape_gradients(co_shape_gradients)
1049 , temp(temp)
1050 , batch_size(batch_size)
1051 {}
1052
1053 template <int dim, int n_rows, int n_columns, typename Number>
1054 template <int direction,
1055 bool dof_to_quad,
1056 bool add,
1057 bool in_place,
1058 typename ViewTypeIn,
1059 typename ViewTypeOut>
1060 DEAL_II_HOST_DEVICE void
1061 EvaluatorTensorProduct<evaluate_general, dim, n_rows, n_columns, Number>::
1062 value(const ViewTypeIn in, ViewTypeOut out) const
1063 {
1064 if constexpr (in_place)
1065 {
1066 apply<dim, direction, n_rows, n_columns, dof_to_quad, false>(
1067 team_member, shape_values, in, temp, batch_size);
1068
1069 constexpr int nn = dof_to_quad ? n_columns : n_rows;
1070 constexpr int n_blocks1 = Utilities::pow(n_columns, direction);
1071 constexpr int n_blocks2 =
1072 Utilities::pow(n_rows, dim - direction - 1);
1073
1074 populate_view<add>(team_member,
1075 out,
1076 temp,
1077 batch_size * n_blocks1 * nn * n_blocks2);
1078 }
1079 else
1080 {
1081 apply<dim, direction, n_rows, n_columns, dof_to_quad, add>(
1082 team_member, shape_values, in, out, batch_size);
1083 }
1084 }
1085
1086 template <int dim, int n_rows, int n_columns, typename Number>
1087 template <int direction,
1088 bool dof_to_quad,
1089 bool add,
1090 bool in_place,
1091 typename ViewTypeIn,
1092 typename ViewTypeOut>
1093 DEAL_II_HOST_DEVICE void
1094 EvaluatorTensorProduct<evaluate_general, dim, n_rows, n_columns, Number>::
1095 derivative(const ViewTypeIn in, ViewTypeOut out) const
1096 {
1097 if constexpr (in_place)
1098 {
1099 apply<dim, direction, n_rows, n_columns, dof_to_quad, false>(
1100 team_member, shape_gradients, in, temp, batch_size);
1101
1102 constexpr int nn = dof_to_quad ? n_columns : n_rows;
1103 constexpr int n_blocks1 = Utilities::pow(n_columns, direction);
1104 constexpr int n_blocks2 =
1105 Utilities::pow(n_rows, dim - direction - 1);
1106
1107 populate_view<add>(team_member,
1108 out,
1109 temp,
1110 batch_size * n_blocks1 * nn * n_blocks2);
1111 }
1112 else
1113 {
1114 apply<dim, direction, n_rows, n_columns, dof_to_quad, add>(
1115 team_member, shape_gradients, in, out, batch_size);
1116 }
1117 }
1118
1119 template <int dim, int n_rows, int n_columns, typename Number>
1120 template <int direction,
1121 bool dof_to_quad,
1122 bool add,
1123 bool in_place,
1124 typename ViewTypeIn,
1125 typename ViewTypeOut>
1126 DEAL_II_HOST_DEVICE void
1127 EvaluatorTensorProduct<evaluate_general, dim, n_rows, n_columns, Number>::
1128 co_derivative(const ViewTypeIn in, ViewTypeOut out) const
1129 {
1130 if constexpr (in_place)
1131 {
1132 apply<dim, direction, n_columns, n_columns, dof_to_quad, false>(
1133 team_member, co_shape_gradients, in, temp, batch_size);
1134
1135 constexpr int n_blocks1 = Utilities::pow(n_columns, direction);
1136 constexpr int n_blocks2 =
1137 Utilities::pow(n_columns, dim - direction - 1);
1138
1139 populate_view<add>(team_member,
1140 out,
1141 temp,
1142 batch_size * n_blocks1 * n_columns * n_blocks2);
1143 }
1144 else
1145 {
1146 apply<dim, direction, n_columns, n_columns, dof_to_quad, add>(
1147 team_member, co_shape_gradients, in, out, batch_size);
1148 }
1149 }
1150
1151 template <int dim, int n_rows, int n_columns, typename Number>
1152 template <bool transpose,
1153 bool add,
1154 typename ViewTypeIn,
1155 typename ViewTypeOut>
1156 DEAL_II_HOST_DEVICE void
1157 EvaluatorTensorProduct<evaluate_general, dim, n_rows, n_columns, Number>::
1158 co_gradient(const ViewTypeIn in, ViewTypeOut out) const
1159 {
1160 static_assert(dim >= 1, "dim must be at least 1");
1161 static_assert(
1162 ViewTypeIn::rank == (transpose ? 2 : 1),
1163 "in must be the values (1D view) for evaluate_gradients (transpose = "
1164 "false), or the gradients (2D view) for integrate_gradients (transpose "
1165 "= true).");
1166 static_assert(
1167 ViewTypeOut::rank == (transpose ? 1 : 2),
1168 "out must be the gradients (2D view) for evaluate_gradients (transpose "
1169 "= false), or the values (1D view) for integrate_gradients (transpose "
1170 "= true).");
1171
1172 constexpr int n_q_points = Utilities::pow(n_columns, dim);
1173 constexpr int co_dimension_size = Utilities::pow(n_columns, dim - 1);
1174
1175 Kokkos::parallel_for(
1176 Kokkos::TeamVectorRange(team_member, batch_size * co_dimension_size),
1177 [&](const int tid) {
1178 const int elmnt_idx = tid / co_dimension_size;
1179 const int reminder = tid % co_dimension_size;
1180
1181 // The main two-nested loop below follows the sum-factorization
1182 // layout of the last (dim-1) direction. We cache in the register
1183 // what varies only in the inner-most loop over 'n', but stays
1184 // constant in the outer loop over 'last', namely co_shape_gradients
1185 // for directions 0,1,..., dim-2, and the in values in the last
1186 // (dim-1) direction.
1187 Kokkos::Array<int, dim - 1> idx_d, stride_d;
1188 Number reg[dim][n_columns];
1189
1190 for (int d = 0; d < dim - 1; ++d)
1191 {
1192 // cache stride and index to access the quad points in 1D view
1193 // layout
1194 stride_d[d] = Utilities::pow(n_columns, d);
1195 idx_d[d] = (reminder / stride_d[d]) % n_columns;
1196
1197 // cache co_shape_gradients in 0,1,..., dim-2 directions
1198 for (int n = 0; n < n_columns; ++n)
1199 {
1200 if constexpr (!transpose)
1201 reg[d][n] = co_shape_gradients(n * n_columns + idx_d[d]);
1202 else
1203 reg[d][n] = co_shape_gradients(idx_d[d] * n_columns + n);
1204 }
1205 }
1206
1207 // cache the in vector values in the last (dim-1) direction
1208 for (int n = 0; n < n_columns; ++n)
1209 {
1210 if constexpr (!transpose)
1211 reg[dim - 1][n] = in(elmnt_idx * n_q_points + reminder +
1212 n * co_dimension_size);
1213 else
1214 reg[dim - 1][n] = in(elmnt_idx * n_q_points + reminder +
1215 n * co_dimension_size,
1216 dim - 1);
1217 }
1218
1219 // Main loop following the sum factotization layout of the last
1220 // (dim-1) direction
1221 for (int last = 0; last < n_columns; ++last)
1222 {
1223 const int q_point = reminder + last * co_dimension_size;
1224
1225 if constexpr (!transpose)
1226 {
1227 Number result[dim];
1228 for (int d = 0; d < dim - 1; ++d)
1229 {
1230 const int q_point_base =
1231 q_point - idx_d[d] * stride_d[d];
1232 const int in_base =
1233 elmnt_idx * n_q_points + q_point_base;
1234
1235 // compute the sum factorization in d<dim-1 direction
1236 // with co_shape_gradients cached
1237 Number res_d = 0;
1238 for (int n = 0; n < n_columns; ++n)
1239 res_d += reg[d][n] * in(in_base + n * stride_d[d]);
1240 result[d] = res_d;
1241 }
1242
1243 // compute the sum factorization in d=dim-1 direction
1244 // with in vector cached
1245 {
1246 Number res_d = 0;
1247 for (int n = 0; n < n_columns; ++n)
1248 res_d += co_shape_gradients(n * n_columns + last) *
1249 reg[dim - 1][n];
1250 result[dim - 1] = res_d;
1251 }
1252
1253 for (int d = 0; d < dim; ++d)
1254 {
1255 if constexpr (add)
1256 out(elmnt_idx * n_q_points + q_point, d) += result[d];
1257 else
1258 out(elmnt_idx * n_q_points + q_point, d) = result[d];
1259 }
1260 }
1261 else
1262 {
1263 Number result = 0;
1264
1265 for (int d = 0; d < dim - 1; ++d)
1266 {
1267 const int point_base = q_point - idx_d[d] * stride_d[d];
1268 const int grad_row =
1269 elmnt_idx * n_q_points + point_base;
1270
1271 // compute the sum factorization in d<dim-1 direction
1272 // with co_shape_gradients cached
1273 for (int n = 0; n < n_columns; ++n)
1274 result +=
1275 in(grad_row + n * stride_d[d], d) * reg[d][n];
1276 }
1277
1278 // compute the sum factorization in d=dim-1 direction
1279 // with in vector cached
1280 for (int n = 0; n < n_columns; ++n)
1281 result += reg[dim - 1][n] *
1282 co_shape_gradients(last * n_columns + n);
1283
1284 if constexpr (add)
1285 out(elmnt_idx * n_q_points + q_point) += result;
1286 else
1287 out(elmnt_idx * n_q_points + q_point) = result;
1288 }
1289 }
1290 });
1291
1292 team_member.team_barrier();
1293 }
1294
1295
1296 } // namespace batched
1297 } // namespace internal
1298} // namespace Portable
1299
1300DEAL_II_NAMESPACE_CLOSE
1301
1302#endif
DEAL_II_NAMESPACE_OPEN
#define DEAL_II_NAMESPACE_OPEN
Definition config.h:38
DEAL_II_HOST_DEVICE
#define DEAL_II_HOST_DEVICE
Definition config.h:171
DEAL_II_NAMESPACE_CLOSE
#define DEAL_II_NAMESPACE_CLOSE
Definition config.h:39
transpose
DerivativeForm< 1, spacedim, dim, Number > transpose(const DerivativeForm< 1, dim, spacedim, Number > &DF)
Definition derivative_form.h:629
Assert
#define Assert(cond, exc)
Definition exception_macros.h:543
StandardExceptions::ExcInternalError
static ::ExceptionBase & ExcInternalError()
memory_space.h
PlasticityLab::Number
*  *  *  RotationFunction< dim, Number >::RotationFunction Number(dim)
Definition ALE_Finite_Strain_Plasticity.h:8733
Portable::internal::batched::populate_view
void populate_view(const Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type &team_member, ViewTypeOut dst, const ViewTypeIn src, const int N)
Definition portable_tensor_product_kernels.h:898
Portable::internal::batched::apply_matrix_vector_product
void apply_matrix_vector_product(const TypeMatrix matrix, const TypeIn in, TypeOut out)
Definition portable_tensor_product_kernels.h:693
Portable::internal::populate_view
void populate_view(const Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type &team_member, ViewTypeOut dst, const ViewTypeIn src, const int N)
Definition portable_tensor_product_kernels.h:52
Portable::internal::apply
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)
Definition portable_tensor_product_kernels.h:300
Utilities::MPI::sum
T sum(const T &t, const MPI_Comm mpi_communicator)
Utilities::pow
constexpr T pow(const T base, const int iexp)
Definition utilities.h:966
MemorySpace::Default::kokkos_space
::Kokkos::DefaultExecutionSpace::memory_space kokkos_space
Definition memory_space.h:43
Portable::internal::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number, ShapeDataMemorySpace >::TeamHandle
Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type TeamHandle
Definition portable_tensor_product_kernels.h:422
Portable::internal::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number, ShapeDataMemorySpace >::SharedView
Kokkos::View< Number *, MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space, Kokkos::MemoryTraits< Kokkos::Unmanaged > > SharedView
Definition portable_tensor_product_kernels.h:427
Portable::internal::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number, ShapeDataMemorySpace >::shape_values
Kokkos::View< Number *, ShapeDataMemorySpace > shape_values
Definition portable_tensor_product_kernels.h:484
Portable::internal::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number, ShapeDataMemorySpace >::team_member
const TeamHandle & team_member
Definition portable_tensor_product_kernels.h:479
Portable::internal::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number, ShapeDataMemorySpace >::co_shape_gradients
Kokkos::View< Number *, ShapeDataMemorySpace > co_shape_gradients
Definition portable_tensor_product_kernels.h:494
Portable::internal::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number, ShapeDataMemorySpace >::shape_gradients
Kokkos::View< Number *, ShapeDataMemorySpace > shape_gradients
Definition portable_tensor_product_kernels.h:489
Portable::internal::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number, ShapeDataMemorySpace >::temp
SharedView temp
Definition portable_tensor_product_kernels.h:499
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::TeamHandle
Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type TeamHandle
Definition portable_tensor_product_kernels.h:943
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::shape_gradients
ShapeDataType shape_gradients
Definition portable_tensor_product_kernels.h:1030
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::team_member
const TeamHandle & team_member
Definition portable_tensor_product_kernels.h:1028
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::ShapeDataType
Kokkos::View< Number *, MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space, Kokkos::MemoryTraits< Kokkos::Unmanaged > > ShapeDataType
Definition portable_tensor_product_kernels.h:949
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::batch_size
const int batch_size
Definition portable_tensor_product_kernels.h:1033
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::temp
SharedView temp
Definition portable_tensor_product_kernels.h:1032
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::shape_values
ShapeDataType shape_values
Definition portable_tensor_product_kernels.h:1029
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::co_shape_gradients
ShapeDataType co_shape_gradients
Definition portable_tensor_product_kernels.h:1031
Portable::internal::batched::EvaluatorTensorProduct< evaluate_general, dim, n_rows, n_columns, Number >::SharedView
Kokkos::View< Number *, MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space, Kokkos::MemoryTraits< Kokkos::Unmanaged > > SharedView
Definition portable_tensor_product_kernels.h:955
Generated by doxygen 1.9.8