14#ifndef dealii__tensor_product_kernels_h
15#define dealii__tensor_product_kernels_h
22#include <Kokkos_Core.hpp>
50 template <
bool add,
typename ViewTypeIn,
typename ViewTypeOut>
53 const Kokkos::TeamPolicy<
54 MemorySpace::Default::kokkos_space::execution_space>::member_type
62 Kokkos::parallel_for(Kokkos::TeamVectorRange(team_member, N),
65 Kokkos::atomic_add(&dst(i), src(i));
70 team_member.team_barrier();
75#if DEAL_II_KOKKOS_VERSION_GTE(4, 0, 0)
83 typename ShapeDataMemorySpace,
84 bool contract_over_rows,
89 apply_1d(
const Kokkos::TeamPolicy<
90 MemorySpace::Default::kokkos_space::execution_space>::member_type
92 const Kokkos::View<Number *, ShapeDataMemorySpace> shape_data,
96 constexpr int Nk = (contract_over_rows ? n_rows : n_columns),
97 Nq = (contract_over_rows ? n_columns : n_rows);
103 Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, Nq),
106 for (
int k = 0; k < Nk; ++k)
108 const int shape_idx =
109 (contract_over_rows ? q + k * Nq :
111 sum += shape_data(shape_idx) * in(k);
115 Kokkos::atomic_add(&out(q), sum);
120 team_member.team_barrier();
128 template <
int n_rows,
132 typename ShapeDataMemorySpace,
133 bool contract_over_rows,
136 typename ViewTypeOut>
138 apply_2d(
const Kokkos::TeamPolicy<
139 MemorySpace::Default::kokkos_space::execution_space>::member_type
141 const Kokkos::View<Number *, ShapeDataMemorySpace> shape_data,
145 using TeamType = Kokkos::TeamPolicy<
146 MemorySpace::Default::kokkos_space::execution_space>::member_type;
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);
174 Kokkos::TeamThreadMDRange<Kokkos::Rank<2>, TeamType>(team_member,
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;
181 const int base_in = (direction == 0 ? j * Nk : j);
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);
189 const int index_out = (direction == 0 ? j * Nq + q : j + q * Nj);
192 Kokkos::atomic_add(&out(index_out), sum);
194 out(index_out) =
sum;
197 team_member.team_barrier();
205 template <
int n_rows,
209 typename ShapeDataMemorySpace,
210 bool contract_over_rows,
213 typename ViewTypeOut>
215 apply_3d(
const Kokkos::TeamPolicy<
216 MemorySpace::Default::kokkos_space::execution_space>::member_type
218 const Kokkos::View<Number *, ShapeDataMemorySpace> shape_data,
222 using TeamType = Kokkos::TeamPolicy<
223 MemorySpace::Default::kokkos_space::execution_space>::member_type;
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);
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;
263 (direction == 0 ? (i * Nj + j) * Nk :
264 (direction == 1 ? i + j * Ni * Nk : i * Nj + j));
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);
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));
278 Kokkos::atomic_add(&out(index_out), sum);
280 out(index_out) =
sum;
283 team_member.team_barrier();
293 typename ShapeDataMemorySpace,
295 bool contract_over_rows,
298 typename ViewTypeOut>
301 MemorySpace::Default::kokkos_space::execution_space>::member_type
303 const Kokkos::View<Number *, ShapeDataMemorySpace> shape_data,
318 if constexpr (dim == 1)
323 ShapeDataMemorySpace,
325 add>(team_member, shape_data, in, out);
326 if constexpr (dim == 2)
331 ShapeDataMemorySpace,
333 add>(team_member, shape_data, in, out);
334 if constexpr (dim == 3)
339 ShapeDataMemorySpace,
341 add>(team_member, shape_data, in, out);
348 constexpr int Nk = contract_over_rows ? n_rows : n_columns;
349 constexpr int Nq = contract_over_rows ? n_columns : n_rows;
355 constexpr int N = NI * NJ * Nq;
358 Kokkos::parallel_for(
359 Kokkos::TeamThreadRange(team_member, N), [&](
const int index_out) {
362 const int q = (index_out / stride) % Nq;
363 const int I = (index_out / stride) / Nq;
364 const int J = index_out % stride;
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;
370 Number sum = shape_data(base_shape) * in(base_in);
371 for (
int k = 1; k < Nk; ++k)
373 const int index_in = (I * Nk + k) * stride + J;
374 sum += shape_data(base_shape + k * stride_shape) * in(index_in);
378 Kokkos::atomic_add(&out(index_out), sum);
380 out(index_out) = sum;
383 team_member.team_barrier();
397 typename ShapeDataMemorySpace =
412 typename ShapeDataMemorySpace>
418 ShapeDataMemorySpace>
422 MemorySpace::Default::kokkos_space::execution_space>::member_type;
425 MemorySpace::Default::kokkos_space::
426 execution_space::scratch_memory_space,
427 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
432 Kokkos::View<Number *, ShapeDataMemorySpace>
shape_values,
441 template <
int direction,
446 typename ViewTypeOut>
448 values(
const ViewTypeIn in, ViewTypeOut out)
const;
454 template <
int direction,
459 typename ViewTypeOut>
461 gradients(
const ViewTypeIn in, ViewTypeOut out)
const;
467 template <
int direction,
472 typename ViewTypeOut>
474 co_gradients(
const ViewTypeIn in, ViewTypeOut out)
const;
508 typename ShapeDataMemorySpace>
515 ShapeDataMemorySpace>::
516 EvaluatorTensorProduct(
518 Kokkos::View<Number *, ShapeDataMemorySpace>
shape_values,
535 typename ShapeDataMemorySpace>
536 template <
int direction,
541 typename ViewTypeOut>
549 ViewTypeOut out)
const
551 if constexpr (in_place)
557 ShapeDataMemorySpace,
569 ShapeDataMemorySpace,
581 typename ShapeDataMemorySpace>
582 template <
int direction,
587 typename ViewTypeOut>
598 if constexpr (in_place)
604 ShapeDataMemorySpace,
616 ShapeDataMemorySpace,
628 typename ShapeDataMemorySpace>
629 template <
int direction,
634 typename ViewTypeOut>
643 ViewTypeOut out)
const
645 if constexpr (in_place)
651 ShapeDataMemorySpace,
663 ShapeDataMemorySpace,
#define DEAL_II_NAMESPACE_OPEN
#define DEAL_II_HOST_DEVICE
#define DEAL_II_NAMESPACE_CLOSE
#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)
::Kokkos::DefaultExecutionSpace::memory_space kokkos_space
Kokkos::TeamPolicy< MemorySpace::Default::kokkos_space::execution_space >::member_type TeamHandle
void gradients(const ViewTypeIn in, ViewTypeOut out) const
void co_gradients(const ViewTypeIn in, ViewTypeOut out) const
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)
Kokkos::View< Number *, ShapeDataMemorySpace > shape_values
const TeamHandle & team_member
Kokkos::View< Number *, ShapeDataMemorySpace > co_shape_gradients
Kokkos::View< Number *, ShapeDataMemorySpace > shape_gradients
void values(const ViewTypeIn in, ViewTypeOut out) const