16#ifndef dealii__tensor_product_kernels_h
17#define dealii__tensor_product_kernels_h
24#include <Kokkos_Core.hpp>
52 template <
bool add,
typename ViewTypeIn,
typename ViewTypeOut>
55 const Kokkos::TeamPolicy<
56 MemorySpace::Default::kokkos_space::execution_space>::member_type
64 Kokkos::parallel_for(Kokkos::TeamVectorRange(team_member, N),
67 Kokkos::atomic_add(&dst(i), src(i));
72 team_member.team_barrier();
77#if DEAL_II_KOKKOS_VERSION_GTE(4, 0, 0)
85 bool contract_over_rows,
90 apply_1d(
const Kokkos::TeamPolicy<
91 MemorySpace::Default::kokkos_space::execution_space>::member_type
93 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
98 constexpr int Nk = (contract_over_rows ? n_rows : n_columns),
99 Nq = (contract_over_rows ? n_columns : n_rows);
105 Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, Nq),
108 for (
int k = 0; k < Nk; ++k)
110 const int shape_idx =
111 (contract_over_rows ? q + k * Nq :
113 sum += shape_data(shape_idx) * in(k);
117 Kokkos::atomic_add(&out(q), sum);
122 team_member.team_barrier();
130 template <
int n_rows,
134 bool contract_over_rows,
137 typename ViewTypeOut>
139 apply_2d(
const Kokkos::TeamPolicy<
140 MemorySpace::Default::kokkos_space::execution_space>::member_type
142 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
147 using TeamType = Kokkos::TeamPolicy<
148 MemorySpace::Default::kokkos_space::execution_space>::member_type;
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);
176 Kokkos::TeamThreadMDRange<Kokkos::Rank<2>, TeamType>(team_member,
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;
183 const int base_in = (direction == 0 ? j * Nk : j);
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);
191 const int index_out = (direction == 0 ? j * Nq + q : j + q * Nj);
194 Kokkos::atomic_add(&out(index_out),
sum);
196 out(index_out) =
sum;
199 team_member.team_barrier();
207 template <
int n_rows,
211 bool contract_over_rows,
214 typename ViewTypeOut>
216 apply_3d(
const Kokkos::TeamPolicy<
217 MemorySpace::Default::kokkos_space::execution_space>::member_type
219 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
224 using TeamType = Kokkos::TeamPolicy<
225 MemorySpace::Default::kokkos_space::execution_space>::member_type;
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);
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;
265 (direction == 0 ? (i * Nj + j) * Nk :
266 (direction == 1 ? i + j * Ni * Nk : i * Nj + j));
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);
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));
280 Kokkos::atomic_add(&out(index_out),
sum);
282 out(index_out) =
sum;
285 team_member.team_barrier();
296 bool contract_over_rows,
299 typename ViewTypeOut>
302 MemorySpace::Default::kokkos_space::execution_space>::member_type
304 const Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
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);
325 constexpr int Nk = contract_over_rows ? n_rows : n_columns;
326 constexpr int Nq = contract_over_rows ? n_columns : n_rows;
332 constexpr int N = NI * NJ * Nq;
335 Kokkos::parallel_for(
336 Kokkos::TeamThreadRange(team_member, N), [&](
const int index_out) {
339 const int q = (index_out / stride) % Nq;
340 const int I = (index_out / stride) / Nq;
341 const int J = index_out % stride;
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;
347 Number sum = shape_data(base_shape) * in(base_in);
348 for (
int k = 1; k < Nk; ++k)
350 const int index_in = (I * Nk + k) * stride + J;
351 sum += shape_data(base_shape + k * stride_shape) * in(index_in);
355 Kokkos::atomic_add(&out(index_out), sum);
357 out(index_out) = sum;
360 team_member.team_barrier();
383 template <
int dim,
int n_rows,
int n_columns,
typename Number>
392 MemorySpace::Default::kokkos_space::execution_space>::member_type;
395 MemorySpace::Default::kokkos_space::
396 execution_space::scratch_memory_space,
397 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
402 Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values,
403 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
405 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
413 template <
int direction,
418 typename ViewTypeOut>
420 values(
const ViewTypeIn in, ViewTypeOut out)
const;
426 template <
int direction,
431 typename ViewTypeOut>
433 gradients(
const ViewTypeIn in, ViewTypeOut out)
const;
439 template <
int direction,
444 typename ViewTypeOut>
446 co_gradients(
const ViewTypeIn in, ViewTypeOut out)
const;
456 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
shape_values;
461 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
467 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
478 template <
int dim,
int n_rows,
int n_columns,
typename Number>
483 Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values,
484 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
486 Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
489 : team_member(team_member)
490 , shape_values(shape_values)
491 , shape_gradients(shape_gradients)
492 , co_shape_gradients(co_shape_gradients)
498 template <
int dim,
int n_rows,
int n_columns,
typename Number>
499 template <
int direction,
504 typename ViewTypeOut>
507 values(
const ViewTypeIn in, ViewTypeOut out)
const
509 if constexpr (in_place)
511 apply<dim, n_rows, n_columns, Number, direction, dof_to_quad, false>(
512 team_member, shape_values, in, temp);
514 populate_view<add>(team_member, out, temp, out.extent(0));
517 apply<dim, n_rows, n_columns, Number, direction, dof_to_quad, add>(
518 team_member, shape_values, in, out);
523 template <
int dim,
int n_rows,
int n_columns,
typename Number>
524 template <
int direction,
529 typename ViewTypeOut>
532 gradients(
const ViewTypeIn in, ViewTypeOut out)
const
534 if constexpr (in_place)
536 apply<dim, n_rows, n_columns, Number, direction, dof_to_quad, false>(
537 team_member, shape_gradients, in, temp);
539 populate_view<add>(team_member, out, temp, out.extent(0));
542 apply<dim, n_rows, n_columns, Number, direction, dof_to_quad, add>(
543 team_member, shape_gradients, in, out);
548 template <
int dim,
int n_rows,
int n_columns,
typename Number>
549 template <
int direction,
554 typename ViewTypeOut>
557 co_gradients(
const ViewTypeIn in, ViewTypeOut out)
const
559 if constexpr (in_place)
567 false>(team_member, co_shape_gradients, in, temp);
569 populate_view<add>(team_member, out, temp, out.extent(0));
572 apply<dim, n_columns, n_columns, Number, direction, dof_to_quad, add>(
573 team_member, co_shape_gradients, in, out);
#define DEAL_II_NAMESPACE_OPEN
#define DEAL_II_HOST_DEVICE
#define DEAL_II_NAMESPACE_CLOSE
#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)
Kokkos::View< Number *, MemorySpace::Default::kokkos_space > shape_values
Kokkos::View< Number *, MemorySpace::Default::kokkos_space > co_shape_gradients
const TeamHandle & team_member
Kokkos::View< Number *, MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space, Kokkos::MemoryTraits< Kokkos::Unmanaged > > SharedView
Kokkos::View< Number *, MemorySpace::Default::kokkos_space > shape_gradients
Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type TeamHandle