bjacobi/bjkokkos/bjkokkoskernels.kokkos.cxx

81     auto d                         = Kokkos::subview(_D, Kokkos::make_pair(first_matrix, last_matrix), Kokkos::ALL);
82     auto x                         = Kokkos::subview(_X, Kokkos::make_pair(first_matrix, last_matrix), Kokkos::ALL);
83     auto b                         = Kokkos::subview(_B, Kokkos::make_pair(first_matrix, last_matrix), Kokkos::ALL);
84     using ScratchPadIntViewType    = Kokkos::View<typename IntView::non_const_value_type *, typename IntView::array_layout, typename IntView::execution_space::scratch_memory_space>;
85     using ScratchPadValuesViewType = Kokkos::View<typename ValuesViewType::non_const_value_type **, typename ValuesViewType::array_layout, typename ValuesViewType::execution_space::scratch_memory_space>;
91     TeamVectorCopy1D::invoke(member, Kokkos::subview(_r, graphID, Kokkos::ALL), r);
92     TeamVectorCopy1D::invoke(member, Kokkos::subview(_c, graphID, Kokkos::ALL), c);
98     KokkosBatched::TeamVectorCopy<MemberType>::invoke(member, Kokkos::subview(_diag, Kokkos::make_pair(first_matrix, last_matrix), Kokkos::ALL), diag);
108     Kokkos::Timer timer;
109     Kokkos::Profiling::pushRegion(name.c_str());
111     Kokkos::TeamPolicy<DeviceType> auto_policy(ceil(1. * _D.extent(0) / _N_team), Kokkos::AUTO(), Kokkos::AUTO());
112     Kokkos::TeamPolicy<DeviceType> tuned_policy(ceil(1. * _D.extent(0) / _N_team), _team_size, _vector_length);
113     Kokkos::TeamPolicy<DeviceType> policy;
130     using ViewType2D    = Kokkos::View<ScalarType **, Layout, EXSP>;
131     using IntViewType1D = Kokkos::View<PetscInt *, Layout, EXSP>;
142     policy.set_scratch_size(0, Kokkos::PerTeam(bytes_tmp));
143     policy.set_scratch_size(1, Kokkos::PerTeam(bytes_col_idc + bytes_row_ptr + bytes_diag));
147     Kokkos::parallel_for(name.c_str(), policy, *this);
171   PetscCheck(jac->const_block_size, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Kokkos (GMRES) solver requires constant block size (but can be made to work with species ordering or N_team==1)");
177   Kokkos::View<PetscScalar **, layout, exec_space, Kokkos::MemoryTraits<Kokkos::Unmanaged>> inv_diag((PetscScalar *)glb_idiag, Nsolves, Nloc); // in correct order
193   Kokkos::deep_copy(batch_x, 0.);
195   Kokkos::parallel_for(
196     "rowOffsets+map", Kokkos::TeamPolicy<>(Nsolves, team_size, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) {
200           Kokkos::parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](const int rowb) { // Nloc
208       Kokkos::parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
213   Kokkos::fence();
215     Kokkos::parallel_for(
216       "prefix sum", Kokkos::TeamPolicy<>(Nsolves / Nsolves_team, 1, 1), KOKKOS_LAMBDA(const team_member team) {
221     Kokkos::fence();
223   Kokkos::parallel_for(
224     "copy matrix", Kokkos::TeamPolicy<>(Nsolves /* /batch_sz */, team_size, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) {
226       Kokkos::parallel_for(Kokkos::TeamThreadRange(team, start, end), [=](const int rowb) {
231         Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, n), [=](const int &i) {
238   Kokkos::fence();
241   using MagnitudeType = typename Kokkos::Details::ArithTraits<ScalarType>::mag_type;
242   //using NormViewType              = Kokkos::View<MagnitudeType *, layout, exec_space>;
243   using Norm2DViewType   = Kokkos::View<MagnitudeType **, layout, exec_space>;
244   using Scalar3DViewType = Kokkos::View<ScalarType ***, layout, exec_space>;
245   using IntViewType      = Kokkos::View<int *, layout, exec_space>;
256   Kokkos::fence();
258   Kokkos::parallel_for(
259     "map", Kokkos::TeamPolicy<>(Nsolves /* /batch_sz */, -1, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) {
262       Kokkos::parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
267   // output assume species major - clone from Kokkos solvers