17 #ifndef dealii_cuda_tensor_product_kernels_h
18 #define dealii_cuda_tensor_product_kernels_h
24 #include <deal.II/matrix_free/cuda_matrix_free.templates.h>
49 #if KOKKOS_VERSION >= 40000
53 template <
int n_q_points_1d,
62 apply_1d(
const Kokkos::TeamPolicy<
63 MemorySpace::Default::kokkos_space::execution_space>::member_type
65 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
70 Number t[n_q_points_1d];
79 for (
int k = 0; k < n_q_points_1d; ++k)
81 const unsigned int shape_idx =
82 dof_to_quad ? (q + k * n_q_points_1d) :
83 (k + q * n_q_points_1d);
84 const unsigned int source_idx = k;
85 t[q] += shape_data[shape_idx] * in(source_idx);
89 if constexpr (in_place)
90 team_member.team_barrier();
94 const unsigned int destination_idx = q;
96 Kokkos::atomic_add(&out(destination_idx), t[q]);
98 out(destination_idx) = t[q];
107 template <
int n_q_points_1d,
114 typename ViewTypeOut>
116 apply_2d(
const Kokkos::TeamPolicy<
117 MemorySpace::Default::kokkos_space::execution_space>::member_type
119 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
124 using TeamType = Kokkos::TeamPolicy<
125 MemorySpace::Default::kokkos_space::execution_space>::member_type;
126 constexpr
unsigned int n_q_points =
Utilities::pow(n_q_points_1d, 2);
128 Number t[n_q_points];
130 Kokkos::TeamThreadMDRange<Kokkos::Rank<2>, TeamType>(team_member,
134 int q_point = i + j * n_q_points_1d;
140 for (
int k = 0; k < n_q_points_1d; ++k)
142 const unsigned int shape_idx =
143 dof_to_quad ? (j + k * n_q_points_1d) : (k + j * n_q_points_1d);
144 const unsigned int source_idx = (direction == 0) ?
145 (k + n_q_points_1d * i) :
146 (i + n_q_points_1d * k);
147 t[q_point] += shape_data[shape_idx] * in(source_idx);
152 team_member.team_barrier();
155 const int q_point = i + j * n_q_points_1d;
156 const unsigned int destination_idx =
157 (direction == 0) ? (j + n_q_points_1d * i) : (i + n_q_points_1d * j);
160 Kokkos::atomic_add(&out(destination_idx), t[q_point]);
162 out(destination_idx) = t[q_point];
171 template <
int n_q_points_1d,
178 typename ViewTypeOut>
180 apply_3d(
const Kokkos::TeamPolicy<
181 MemorySpace::Default::kokkos_space::execution_space>::member_type
183 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
188 using TeamType = Kokkos::TeamPolicy<
189 MemorySpace::Default::kokkos_space::execution_space>::member_type;
190 constexpr
unsigned int n_q_points =
Utilities::pow(n_q_points_1d, 3);
192 Number t[n_q_points];
193 auto thread_policy = Kokkos::TeamThreadMDRange<Kokkos::Rank<3>, TeamType>(
194 team_member, n_q_points_1d, n_q_points_1d, n_q_points_1d);
196 thread_policy, [&](
const int i,
const int j,
const int q) {
198 i + j * n_q_points_1d + q * n_q_points_1d * n_q_points_1d;
204 for (
int k = 0; k < n_q_points_1d; ++k)
206 const unsigned int shape_idx =
207 dof_to_quad ? (q + k * n_q_points_1d) : (k + q * n_q_points_1d);
208 const unsigned int source_idx =
210 (k + n_q_points_1d * (i + n_q_points_1d * j)) :
212 (i + n_q_points_1d * (k + n_q_points_1d * j)) :
213 (i + n_q_points_1d * (j + n_q_points_1d * k));
214 t[q_point] += shape_data[shape_idx] * in(source_idx);
219 team_member.team_barrier();
222 thread_policy, [&](
const int i,
const int j,
const int q) {
224 i + j * n_q_points_1d + q * n_q_points_1d * n_q_points_1d;
225 const unsigned int destination_idx =
226 (direction == 0) ? (q + n_q_points_1d * (i + n_q_points_1d * j)) :
227 (direction == 1) ? (i + n_q_points_1d * (q + n_q_points_1d * j)) :
228 (i + n_q_points_1d * (j + n_q_points_1d * q));
231 Kokkos::atomic_add(&out(destination_idx), t[q_point]);
233 out(destination_idx) = t[q_point];
251 typename ViewTypeOut>
254 MemorySpace::Default::kokkos_space::execution_space>::member_type
256 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
261 #if KOKKOS_VERSION >= 40000
262 if constexpr (dim == 1)
263 apply_1d<n_q_points_1d, Number, direction, dof_to_quad, add, in_place>(
264 team_member, shape_data, in, out);
265 if constexpr (dim == 2)
266 apply_2d<n_q_points_1d, Number, direction, dof_to_quad, add, in_place>(
267 team_member, shape_data, in, out);
268 if constexpr (dim == 3)
269 apply_3d<n_q_points_1d, Number, direction, dof_to_quad, add, in_place>(
270 team_member, shape_data, in, out);
272 constexpr
unsigned int n_q_points =
Utilities::pow(n_q_points_1d, dim);
274 Number t[n_q_points];
276 Kokkos::TeamThreadRange(team_member, n_q_points),
277 [&](
const int &q_point) {
278 const unsigned int i = (dim == 1) ? 0 : q_point % n_q_points_1d;
279 const unsigned int j =
280 (dim == 3) ? (q_point / n_q_points_1d) % n_q_points_1d : 0;
281 const unsigned int q =
282 (dim == 1) ? q_point :
283 (dim == 2) ? (q_point / n_q_points_1d) % n_q_points_1d :
284 q_point / (n_q_points_1d * n_q_points_1d);
289 for (
int k = 0; k < n_q_points_1d; ++k)
291 const unsigned int shape_idx =
292 dof_to_quad ? (q + k * n_q_points_1d) : (k + q * n_q_points_1d);
293 const unsigned int source_idx =
295 (k + n_q_points_1d * (i + n_q_points_1d * j)) :
297 (i + n_q_points_1d * (k + n_q_points_1d * j)) :
298 (i + n_q_points_1d * (j + n_q_points_1d * k));
299 t[q_point] += shape_data[shape_idx] *
300 (in_place ? out(source_idx) : in(source_idx));
305 team_member.team_barrier();
308 Kokkos::TeamThreadRange(team_member, n_q_points),
309 [&](
const int &q_point) {
310 const unsigned int i = (dim == 1) ? 0 : q_point % n_q_points_1d;
311 const unsigned int j =
312 (dim == 3) ? (q_point / n_q_points_1d) % n_q_points_1d : 0;
313 const unsigned int q =
314 (dim == 1) ? q_point :
315 (dim == 2) ? (q_point / n_q_points_1d) % n_q_points_1d :
316 q_point / (n_q_points_1d * n_q_points_1d);
318 const unsigned int destination_idx =
319 (direction == 0) ? (q + n_q_points_1d * (i + n_q_points_1d * j)) :
320 (direction == 1) ? (i + n_q_points_1d * (q + n_q_points_1d * j)) :
321 (i + n_q_points_1d * (j + n_q_points_1d * q));
324 Kokkos::atomic_add(&out(destination_idx), t[q_point]);
326 out(destination_idx) = t[q_point];
353 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
361 MemorySpace::Default::kokkos_space::execution_space>::member_type;
366 Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values,
367 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
369 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
376 template <
int direction,
381 typename ViewTypeOut>
383 values(
const ViewTypeIn in, ViewTypeOut out)
const;
389 template <
int direction,
394 typename ViewTypeOut>
396 gradients(
const ViewTypeIn in, ViewTypeOut out)
const;
402 template <
int direction,
407 typename ViewTypeOut>
409 co_gradients(
const ViewTypeIn in, ViewTypeOut out)
const;
414 template <
typename ViewType>
416 value_at_quad_pts(ViewType u);
421 template <
typename ViewType>
423 integrate_value(ViewType u);
429 template <
typename ViewTypeIn,
typename ViewTypeOut>
431 gradient_at_quad_pts(
const ViewTypeIn u, ViewTypeOut grad_u);
437 template <
typename ViewType1,
typename ViewType2>
439 value_and_gradient_at_quad_pts(ViewType1 u, ViewType2 grad_u);
445 template <
bool add,
typename ViewType1,
typename ViewType2>
447 integrate_gradient(ViewType1 u, ViewType2 grad_u);
453 template <
typename ViewType1,
typename ViewType2>
455 integrate_value_and_gradient(ViewType1 u, ViewType2 grad_u);
465 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
shape_values;
470 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
476 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
482 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
489 EvaluatorTensorProduct(
491 Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values,
492 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
494 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
496 : team_member(team_member)
497 , shape_values(shape_values)
498 , shape_gradients(shape_gradients)
499 , co_shape_gradients(co_shape_gradients)
504 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
505 template <
int direction,
510 typename ViewTypeOut>
517 ViewTypeOut out)
const
519 apply<dim, n_q_points_1d, Number, direction, dof_to_quad, add, in_place>(
520 team_member, shape_values, in, out);
525 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
526 template <
int direction,
531 typename ViewTypeOut>
538 ViewTypeOut out)
const
540 apply<dim, n_q_points_1d, Number, direction, dof_to_quad, add, in_place>(
541 team_member, shape_gradients, in, out);
546 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
547 template <
int direction,
552 typename ViewTypeOut>
558 Number>::co_gradients(
const ViewTypeIn in,
559 ViewTypeOut out)
const
561 apply<dim, n_q_points_1d, Number, direction, dof_to_quad, add, in_place>(
562 team_member, co_shape_gradients, in, out);
567 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
568 template <
typename ViewType>
574 Number>::value_at_quad_pts(ViewType u)
576 if constexpr (dim == 1)
577 values<0, true, false, true>(u, u);
578 else if constexpr (dim == 2)
580 values<0, true, false, true>(u, u);
581 team_member.team_barrier();
582 values<1, true, false, true>(u, u);
584 else if constexpr (dim == 3)
586 values<0, true, false, true>(u, u);
587 team_member.team_barrier();
588 values<1, true, false, true>(u, u);
589 team_member.team_barrier();
590 values<2, true, false, true>(u, u);
598 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
599 template <
typename ViewType>
605 Number>::integrate_value(ViewType u)
607 if constexpr (dim == 1)
608 values<0, false, false, true>(u, u);
609 else if constexpr (dim == 2)
611 values<0, false, false, true>(u, u);
612 team_member.team_barrier();
613 values<1, false, false, true>(u, u);
615 else if constexpr (dim == 3)
617 values<0, false, false, true>(u, u);
618 team_member.team_barrier();
619 values<1, false, false, true>(u, u);
620 team_member.team_barrier();
621 values<2, false, false, true>(u, u);
629 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
630 template <
typename ViewTypeIn,
typename ViewTypeOut>
636 Number>::gradient_at_quad_pts(
const ViewTypeIn u,
639 if constexpr (dim == 1)
641 gradients<0, true, false, false>(
642 u, Kokkos::subview(grad_u, Kokkos::ALL, 0));
644 else if constexpr (dim == 2)
646 gradients<0, true, false, false>(
647 u, Kokkos::subview(grad_u, Kokkos::ALL, 0));
648 values<0, true, false, false>(
649 u, Kokkos::subview(grad_u, Kokkos::ALL, 1));
651 team_member.team_barrier();
653 values<1, true, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 0),
654 Kokkos::subview(grad_u, Kokkos::ALL, 0));
655 gradients<1, true, false, true>(
656 Kokkos::subview(grad_u, Kokkos::ALL, 1),
657 Kokkos::subview(grad_u, Kokkos::ALL, 1));
659 else if constexpr (dim == 3)
661 gradients<0, true, false, false>(
662 u, Kokkos::subview(grad_u, Kokkos::ALL, 0));
663 values<0, true, false, false>(
664 u, Kokkos::subview(grad_u, Kokkos::ALL, 1));
665 values<0, true, false, false>(
666 u, Kokkos::subview(grad_u, Kokkos::ALL, 2));
668 team_member.team_barrier();
670 values<1, true, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 0),
671 Kokkos::subview(grad_u, Kokkos::ALL, 0));
672 gradients<1, true, false, true>(
673 Kokkos::subview(grad_u, Kokkos::ALL, 1),
674 Kokkos::subview(grad_u, Kokkos::ALL, 1));
675 values<1, true, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 2),
676 Kokkos::subview(grad_u, Kokkos::ALL, 2));
678 team_member.team_barrier();
680 values<2, true, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 0),
681 Kokkos::subview(grad_u, Kokkos::ALL, 0));
682 values<2, true, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 1),
683 Kokkos::subview(grad_u, Kokkos::ALL, 1));
684 gradients<2, true, false, true>(
685 Kokkos::subview(grad_u, Kokkos::ALL, 2),
686 Kokkos::subview(grad_u, Kokkos::ALL, 2));
694 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
695 template <
typename ViewType1,
typename ViewType2>
701 Number>::value_and_gradient_at_quad_pts(ViewType1 u,
705 if constexpr (dim == 1)
707 values<0, true, false, true>(u, u);
708 team_member.team_barrier();
710 co_gradients<0, true, false, false>(
711 u, Kokkos::subview(grad_u, Kokkos::ALL, 0));
713 else if constexpr (dim == 2)
715 values<0, true, false, true>(u, u);
716 team_member.team_barrier();
717 values<1, true, false, true>(u, u);
718 team_member.team_barrier();
720 co_gradients<0, true, false, false>(
721 u, Kokkos::subview(grad_u, Kokkos::ALL, 0));
722 co_gradients<1, true, false, false>(
723 u, Kokkos::subview(grad_u, Kokkos::ALL, 1));
725 else if constexpr (dim == 3)
727 values<0, true, false, true>(u, u);
728 team_member.team_barrier();
729 values<1, true, false, true>(u, u);
730 team_member.team_barrier();
731 values<2, true, false, true>(u, u);
732 team_member.team_barrier();
734 co_gradients<0, true, false, false>(
735 u, Kokkos::subview(grad_u, Kokkos::ALL, 0));
736 co_gradients<1, true, false, false>(
737 u, Kokkos::subview(grad_u, Kokkos::ALL, 1));
738 co_gradients<2, true, false, false>(
739 u, Kokkos::subview(grad_u, Kokkos::ALL, 2));
747 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
748 template <
bool add,
typename ViewType1,
typename ViewType2>
754 Number>::integrate_gradient(ViewType1 u,
757 if constexpr (dim == 1)
759 gradients<0, false, add, false>(
760 Kokkos::subview(grad_u, Kokkos::ALL, dim), u);
762 else if constexpr (dim == 2)
764 gradients<0, false, false, true>(
765 Kokkos::subview(grad_u, Kokkos::ALL, 0),
766 Kokkos::subview(grad_u, Kokkos::ALL, 0));
767 values<0, false, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 1),
768 Kokkos::subview(grad_u,
772 team_member.team_barrier();
774 values<1, false, add, false>(Kokkos::subview(grad_u, Kokkos::ALL, 0),
776 team_member.team_barrier();
777 gradients<1, false, true, false>(
778 Kokkos::subview(grad_u, Kokkos::ALL, 1), u);
780 else if constexpr (dim == 3)
782 gradients<0, false, false, true>(
783 Kokkos::subview(grad_u, Kokkos::ALL, 0),
784 Kokkos::subview(grad_u, Kokkos::ALL, 0));
785 values<0, false, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 1),
786 Kokkos::subview(grad_u,
789 values<0, false, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 2),
790 Kokkos::subview(grad_u,
794 team_member.team_barrier();
796 values<1, false, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 0),
797 Kokkos::subview(grad_u,
800 gradients<1, false, false, true>(
801 Kokkos::subview(grad_u, Kokkos::ALL, 1),
802 Kokkos::subview(grad_u, Kokkos::ALL, 1));
803 values<1, false, false, true>(Kokkos::subview(grad_u, Kokkos::ALL, 2),
804 Kokkos::subview(grad_u,
808 team_member.team_barrier();
810 values<2, false, add, false>(Kokkos::subview(grad_u, Kokkos::ALL, 0),
812 team_member.team_barrier();
813 values<2, false, true, false>(Kokkos::subview(grad_u, Kokkos::ALL, 1),
815 team_member.team_barrier();
816 gradients<2, false, true, false>(
817 Kokkos::subview(grad_u, Kokkos::ALL, 2), u);
825 template <
int dim,
int fe_degree,
int n_q_po
ints_1d,
typename Number>
826 template <
typename ViewType1,
typename ViewType2>
832 Number>::integrate_value_and_gradient(ViewType1 u,
836 if constexpr (dim == 1)
838 co_gradients<0, false, true, false>(
839 Kokkos::subview(grad_u, Kokkos::ALL, 0), u);
840 team_member.team_barrier();
842 values<0, false, false, true>(u, u);
844 else if constexpr (dim == 2)
846 co_gradients<1, false, true, false>(
847 Kokkos::subview(grad_u, Kokkos::ALL, 1), u);
848 team_member.team_barrier();
849 co_gradients<0, false, true, false>(
850 Kokkos::subview(grad_u, Kokkos::ALL, 0), u);
851 team_member.team_barrier();
853 values<1, false, false, true>(u, u);
854 team_member.team_barrier();
855 values<0, false, false, true>(u, u);
856 team_member.team_barrier();
858 else if constexpr (dim == 3)
860 co_gradients<2, false, true, false>(
861 Kokkos::subview(grad_u, Kokkos::ALL, 2), u);
862 team_member.team_barrier();
863 co_gradients<1, false, true, false>(
864 Kokkos::subview(grad_u, Kokkos::ALL, 1), u);
865 team_member.team_barrier();
866 co_gradients<0, false, true, false>(
867 Kokkos::subview(grad_u, Kokkos::ALL, 0), u);
868 team_member.team_barrier();
870 values<2, false, false, true>(u, u);
871 team_member.team_barrier();
872 values<1, false, false, true>(u, u);
873 team_member.team_barrier();
874 values<0, false, false, true>(u, u);
875 team_member.team_barrier();
#define DEAL_II_NAMESPACE_OPEN
#define DEAL_II_NAMESPACE_CLOSE
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)
constexpr T pow(const T base, const int iexp)
void abort(const ExceptionBase &exc) noexcept
void parallel_for(Iterator x_begin, Iterator x_end, const Functor &functor, const unsigned int grainsize)
#define DEAL_II_HOST_DEVICE
Kokkos::View< Number *, MemorySpace::Default::kokkos_space > shape_gradients
Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type TeamHandle
Kokkos::View< Number *, MemorySpace::Default::kokkos_space > shape_values
const TeamHandle & team_member
Kokkos::View< Number *, MemorySpace::Default::kokkos_space > co_shape_gradients