1 #include <petscvec_kokkos.hpp> 2 #include <petsc/private/pcimpl.h> 3 #include <petsc/private/kspimpl.h> 4 #include <petscksp.h> /*I "petscksp.h" I*/ 5 #include "petscsection.h" 6 #include <petscdmcomposite.h> 7 #include "Kokkos_Core.hpp" 8 9 #include <../src/mat/impls/aij/seq/aij.h> 10 #include <../src/mat/impls/aij/seq/kokkos/aijkok.hpp> 11 12 #if defined(PETSC_HAVE_CUDA_NVTX) 13 #include <nvToolsExt.h> 14 #endif 15 16 #define PCBJKOKKOS_SHARED_LEVEL 0 // 0 is shared, 1 is global 17 #define PCBJKOKKOS_VEC_SIZE 16 18 #define PCBJKOKKOS_TEAM_SIZE 16 19 20 #define PCBJKOKKOS_VERBOSE_LEVEL 2 21 22 typedef Kokkos::DefaultExecutionSpace exec_space; 23 using layout = Kokkos::LayoutRight; 24 using IntView = Kokkos::View<PetscInt **, layout, exec_space>; 25 using AMatrixValueView = const Kokkos::View<PetscScalar **, layout, exec_space>; 26 using XYType = const Kokkos::View<PetscScalar **, layout, exec_space>; 27 28 typedef enum { 29 BATCH_KSP_BICG_IDX, 30 BATCH_KSP_TFQMR_IDX, 31 BATCH_KSP_GMRES_IDX, 32 NUM_BATCH_TYPES 33 } KSPIndex; 34 typedef struct { 35 Vec vec_diag; 36 PetscInt nBlocks; /* total number of blocks */ 37 PetscInt n; // cache host version of d_bid_eqOffset_k[nBlocks] 38 KSP ksp; // Used just for options. Should have one for each block 39 Kokkos::View<PetscInt *, Kokkos::LayoutRight> *d_bid_eqOffset_k; 40 Kokkos::View<PetscScalar *, Kokkos::LayoutRight> *d_idiag_k; 41 Kokkos::View<PetscInt *> *d_isrow_k; 42 Kokkos::View<PetscInt *> *d_isicol_k; 43 KSPIndex ksp_type_idx; 44 PetscInt nwork; 45 PetscInt const_block_size; // used to decide to use shared memory for work vectors 46 PetscInt *dm_Nf; // Number of fields in each DM 47 PetscInt num_dms; 48 // diagnostics 49 PetscBool reason; 50 PetscBool monitor; 51 PetscInt batch_target; 52 PetscInt nsolves_team; 53 PetscInt max_nits; 54 // caches 55 IntView *rowOffsets; 56 IntView *colIndices; 57 XYType *batch_b; 58 XYType *batch_x; 59 AMatrixValueView *batch_values; 60 } PC_PCBJKOKKOS; 61 62 #if defined(PETSC_HAVE_KOKKOS_KERNELS_GMRES) 63 #include <fstream> 64 65 #include "Kokkos_Timer.hpp" 66 #include "Kokkos_Random.hpp" 67 #include "Kokkos_UnorderedMap.hpp" 68 #include "Kokkos_Sort.hpp" 69 70 /// KokkosKernels headers 71 #include "KokkosBatched_Util.hpp" 72 #include "KokkosBatched_Vector.hpp" 73 74 #include <Kokkos_ArithTraits.hpp> 75 #include <KokkosBatched_Util.hpp> 76 #include <KokkosBatched_Vector.hpp> 77 #include <KokkosBatched_Copy_Decl.hpp> 78 #include <KokkosBatched_Copy_Impl.hpp> 79 #include <KokkosBatched_AddRadial_Decl.hpp> 80 #include <KokkosBatched_AddRadial_Impl.hpp> 81 #include <KokkosBatched_Gemm_Decl.hpp> 82 #include <KokkosBatched_Gemm_Serial_Impl.hpp> 83 #include <KokkosBatched_Gemm_Team_Impl.hpp> 84 #include <KokkosBatched_Gemv_Decl.hpp> 85 #include <KokkosBatched_Gemv_Serial_Impl.hpp> 86 #include <KokkosBatched_Gemv_Team_Impl.hpp> 87 #include <KokkosBatched_Trsm_Decl.hpp> 88 #include <KokkosBatched_Trsm_Serial_Impl.hpp> 89 #include <KokkosBatched_Trsm_Team_Impl.hpp> 90 #include <KokkosBatched_Trsv_Decl.hpp> 91 #include <KokkosBatched_Trsv_Serial_Impl.hpp> 92 #include <KokkosBatched_Trsv_Team_Impl.hpp> 93 #include <KokkosBatched_LU_Decl.hpp> 94 #include <KokkosBatched_LU_Serial_Impl.hpp> 95 #include <KokkosBatched_LU_Team_Impl.hpp> 96 #include <KokkosSparse_CrsMatrix.hpp> 97 #include "KokkosBatched_Spmv.hpp" 98 #include "KokkosBatched_CrsMatrix.hpp" 99 #include "KokkosBatched_Krylov_Handle.hpp" 100 #include "KokkosBatched_GMRES.hpp" 101 #include "KokkosBatched_JacobiPrec.hpp" 102 103 template <typename DeviceType, typename ValuesViewType, typename IntView, typename VectorViewType, typename KrylovHandleType> 104 struct Functor_TestBatchedTeamVectorGMRES { 105 const ValuesViewType _D; 106 const ValuesViewType _diag; 107 const IntView _r; 108 const IntView _c; 109 const VectorViewType _X; 110 const VectorViewType _B; 111 const int _N_team, _team_size, _vector_length; 112 const int _N_iteration; 113 const double _tol; 114 const int _ortho_strategy; 115 const int _scratch_pad_level; 116 KrylovHandleType _handle; 117 118 KOKKOS_INLINE_FUNCTION 119 Functor_TestBatchedTeamVectorGMRES(const ValuesViewType &D, const IntView &r, const IntView &c, const VectorViewType &X, const VectorViewType &B, const int N_team, const int team_size, const int vector_length, const int N_iteration, const double tol, const int ortho_strategy, const int scratch_pad_level, KrylovHandleType &handle) : 120 _D(D), _r(r), _c(c), _X(X), _B(B), _N_team(N_team), _team_size(team_size), _vector_length(vector_length), _N_iteration(N_iteration), _tol(tol), _ortho_strategy(ortho_strategy), _scratch_pad_level(scratch_pad_level), _handle(handle) { } 121 122 KOKKOS_INLINE_FUNCTION 123 Functor_TestBatchedTeamVectorGMRES(const ValuesViewType &D, const ValuesViewType &diag, const IntView &r, const IntView &c, const VectorViewType &X, const VectorViewType &B, const int N_team, const int team_size, const int vector_length, const int N_iteration, const double tol, int ortho_strategy, const int scratch_pad_level, KrylovHandleType &handle) : 124 _D(D), _diag(diag), _r(r), _c(c), _X(X), _B(B), _N_team(N_team), _team_size(team_size), _vector_length(vector_length), _N_iteration(N_iteration), _tol(tol), _ortho_strategy(ortho_strategy), _scratch_pad_level(scratch_pad_level), _handle(handle) { } 125 126 template <typename MemberType> 127 KOKKOS_INLINE_FUNCTION void operator()(const MemberType &member) const { 128 const int first_matrix = static_cast<int>(member.league_rank()) * _N_team; 129 const int N = _D.extent(0); 130 const int last_matrix = (static_cast<int>(member.league_rank() + 1) * _N_team < N ? static_cast<int>(member.league_rank() + 1) * _N_team : N); 131 const int graphID = static_cast<int>(member.league_rank()); 132 using TeamVectorCopy1D = KokkosBatched::TeamVectorCopy<MemberType, KokkosBatched::Trans::NoTranspose, 1>; 133 134 auto d = Kokkos::subview(_D, Kokkos::make_pair(first_matrix, last_matrix), Kokkos::ALL); 135 auto x = Kokkos::subview(_X, Kokkos::make_pair(first_matrix, last_matrix), Kokkos::ALL); 136 auto b = Kokkos::subview(_B, Kokkos::make_pair(first_matrix, last_matrix), Kokkos::ALL); 137 using ScratchPadIntViewType = Kokkos::View<typename IntView::non_const_value_type *, typename IntView::array_layout, typename IntView::execution_space::scratch_memory_space>; 138 using ScratchPadValuesViewType = Kokkos::View<typename ValuesViewType::non_const_value_type **, typename ValuesViewType::array_layout, typename ValuesViewType::execution_space::scratch_memory_space>; 139 140 using Operator = KokkosBatched::CrsMatrix<ValuesViewType, ScratchPadIntViewType>; 141 ScratchPadIntViewType r(member.team_scratch(1), _r.extent(1)); 142 ScratchPadIntViewType c(member.team_scratch(1), _c.extent(1)); 143 144 TeamVectorCopy1D::invoke(member, Kokkos::subview(_r, graphID, Kokkos::ALL), r); 145 TeamVectorCopy1D::invoke(member, Kokkos::subview(_c, graphID, Kokkos::ALL), c); 146 Operator A(d, r, c); 147 148 ScratchPadValuesViewType diag(member.team_scratch(1), last_matrix - first_matrix, _diag.extent(1)); 149 using PrecOperator = KokkosBatched::JacobiPrec<ScratchPadValuesViewType>; 150 151 KokkosBatched::TeamVectorCopy<MemberType>::invoke(member, Kokkos::subview(_diag, Kokkos::make_pair(first_matrix, last_matrix), Kokkos::ALL), diag); 152 PrecOperator P(diag); 153 P.setComputedInverse(); 154 155 KokkosBatched::TeamVectorGMRES<MemberType>::template invoke<Operator, VectorViewType, PrecOperator, KrylovHandleType>(member, A, b, x, P, _handle); 156 } 157 inline double run(PC pc) { 158 typedef typename ValuesViewType::value_type value_type; 159 std::string name("KokkosBatched::Test::TeamVectorGMRES"); 160 Kokkos::Timer timer; 161 Kokkos::Profiling::pushRegion(name.c_str()); 162 163 Kokkos::TeamPolicy<DeviceType> auto_policy(ceil(1. * _D.extent(0) / _N_team), Kokkos::AUTO(), Kokkos::AUTO()); 164 Kokkos::TeamPolicy<DeviceType> tuned_policy(ceil(1. * _D.extent(0) / _N_team), _team_size, _vector_length); 165 Kokkos::TeamPolicy<DeviceType> policy; 166 167 if (_team_size < 1) policy = auto_policy; 168 else policy = tuned_policy; 169 170 _handle.set_max_iteration(_N_iteration); 171 _handle.set_tolerance(_tol); 172 _handle.set_ortho_strategy(_ortho_strategy); 173 _handle.set_scratch_pad_level(_scratch_pad_level); 174 _handle.set_compute_last_residual(true); 175 176 int maximum_iteration = _handle.get_max_iteration(); 177 178 using ScalarType = typename ValuesViewType::non_const_value_type; 179 using Layout = typename ValuesViewType::array_layout; 180 using EXSP = typename ValuesViewType::execution_space; 181 182 using MagnitudeType = typename Kokkos::Details::ArithTraits<ScalarType>::mag_type; 183 184 using ViewType1D = Kokkos::View<MagnitudeType *, Layout, EXSP>; 185 using ViewType2D = Kokkos::View<ScalarType **, Layout, EXSP>; 186 using ViewType3D = Kokkos::View<ScalarType ***, Layout, EXSP>; 187 using IntViewType1D = Kokkos::View<PetscInt *, Layout, EXSP>; 188 189 size_t bytes_1D = ViewType2D::shmem_size(_N_team, 1); 190 size_t bytes_row_ptr = IntViewType1D::shmem_size(_r.extent(1)); 191 size_t bytes_col_idc = IntViewType1D::shmem_size(_c.extent(1)); 192 size_t bytes_2D_1 = ViewType2D::shmem_size(_N_team, _X.extent(1)); 193 size_t bytes_2D_2 = ViewType2D::shmem_size(_N_team, maximum_iteration + 1); 194 195 size_t bytes_diag = bytes_2D_1; 196 size_t bytes_tmp = 2 * bytes_2D_1 + 2 * bytes_1D + bytes_2D_2; 197 198 policy.set_scratch_size(0, Kokkos::PerTeam(bytes_tmp)); 199 policy.set_scratch_size(1, Kokkos::PerTeam(bytes_col_idc + bytes_row_ptr + bytes_diag)); 200 PetscInfo(pc, "%d scratch memory(0) = %d + %d + %d bytes_diag=%d; %d scratch memory(1); %d maximum_iterations\n", (int)(bytes_tmp), 2 * (int)bytes_2D_1, 2 * (int)bytes_1D, (int)bytes_2D_2, (int)bytes_diag, (int)(bytes_row_ptr + bytes_col_idc + bytes_diag), (int)maximum_iteration); 201 exec_space().fence(); 202 timer.reset(); 203 Kokkos::parallel_for(name.c_str(), policy, *this); 204 exec_space().fence(); 205 double sec = timer.seconds(); 206 207 return sec; 208 } 209 }; 210 #endif // KK GMRES 211 212 typedef Kokkos::TeamPolicy<>::member_type team_member; 213 214 static PetscErrorCode PCBJKOKKOSCreateKSP_BJKOKKOS(PC pc) { 215 const char *prefix; 216 PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data; 217 DM dm; 218 219 PetscFunctionBegin; 220 PetscCall(KSPCreate(PetscObjectComm((PetscObject)pc), &jac->ksp)); 221 PetscCall(KSPSetErrorIfNotConverged(jac->ksp, pc->erroriffailure)); 222 PetscCall(PetscObjectIncrementTabLevel((PetscObject)jac->ksp, (PetscObject)pc, 1)); 223 PetscCall(PCGetOptionsPrefix(pc, &prefix)); 224 PetscCall(KSPSetOptionsPrefix(jac->ksp, prefix)); 225 PetscCall(KSPAppendOptionsPrefix(jac->ksp, "pc_bjkokkos_")); 226 PetscCall(PCGetDM(pc, &dm)); 227 if (dm) { 228 PetscCall(KSPSetDM(jac->ksp, dm)); 229 PetscCall(KSPSetDMActive(jac->ksp, PETSC_FALSE)); 230 } 231 jac->reason = PETSC_FALSE; 232 jac->monitor = PETSC_FALSE; 233 jac->batch_target = -1; 234 jac->nsolves_team = 1; 235 jac->ksp->max_it = 50; // this is realy for GMRES w/o restarts 236 PetscFunctionReturn(0); 237 } 238 239 // y <-- Ax 240 KOKKOS_INLINE_FUNCTION PetscErrorCode MatMult(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, const PetscInt start, const PetscInt end, const PetscScalar *x_loc, PetscScalar *y_loc) { 241 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, start, end), [=](const int rowb) { 242 int rowa = ic[rowb]; 243 int n = glb_Aai[rowa + 1] - glb_Aai[rowa]; 244 const PetscInt *aj = glb_Aaj + glb_Aai[rowa]; 245 const PetscScalar *aa = glb_Aaa + glb_Aai[rowa]; 246 PetscScalar sum; 247 Kokkos::parallel_reduce( 248 Kokkos::ThreadVectorRange(team, n), [=](const int i, PetscScalar &lsum) { lsum += aa[i] * x_loc[r[aj[i]] - start]; }, sum); 249 Kokkos::single(Kokkos::PerThread(team), [=]() { y_loc[rowb - start] = sum; }); 250 }); 251 team.team_barrier(); 252 return 0; 253 } 254 255 // temp buffer per thread with reduction at end? 256 KOKKOS_INLINE_FUNCTION PetscErrorCode MatMultTranspose(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, const PetscInt start, const PetscInt end, const PetscScalar *x_loc, PetscScalar *y_loc) { 257 Kokkos::parallel_for(Kokkos::TeamVectorRange(team, end - start), [=](int i) { y_loc[i] = 0; }); 258 team.team_barrier(); 259 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, start, end), [=](const int rowb) { 260 int rowa = ic[rowb]; 261 int n = glb_Aai[rowa + 1] - glb_Aai[rowa]; 262 const PetscInt *aj = glb_Aaj + glb_Aai[rowa]; 263 const PetscScalar *aa = glb_Aaa + glb_Aai[rowa]; 264 const PetscScalar xx = x_loc[rowb - start]; // rowb = ic[rowa] = ic[r[rowb]] 265 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, n), [=](const int &i) { 266 PetscScalar val = aa[i] * xx; 267 Kokkos::atomic_fetch_add(&y_loc[r[aj[i]] - start], val); 268 }); 269 }); 270 team.team_barrier(); 271 return 0; 272 } 273 274 typedef struct Batch_MetaData_TAG { 275 PetscInt flops; 276 PetscInt its; 277 KSPConvergedReason reason; 278 } Batch_MetaData; 279 280 // Solve A(BB^-1)x = y with TFQMR. Right preconditioned to get un-preconditioned residual 281 KOKKOS_INLINE_FUNCTION PetscErrorCode BJSolve_TFQMR(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, PetscScalar *work_space_global, const int stride_global, const int nShareVec, PetscScalar *work_space_shared, const int stride_shared, PetscReal rtol, PetscReal atol, PetscReal dtol, PetscInt maxit, Batch_MetaData *metad, const PetscInt start, const PetscInt end, const PetscScalar glb_idiag[], const PetscScalar *glb_b, PetscScalar *glb_x, bool monitor) { 282 using Kokkos::parallel_for; 283 using Kokkos::parallel_reduce; 284 int Nblk = end - start, i, m, stride = stride_shared, idx = 0; 285 PetscReal dp, dpold, w, dpest, tau, psi, cm, r0; 286 const PetscScalar *Diag = &glb_idiag[start]; 287 PetscScalar *ptr = work_space_shared, rho, rhoold, a, s, b, eta, etaold, psiold, cf, dpi; 288 289 if (idx++ == nShareVec) { 290 ptr = work_space_global; 291 stride = stride_global; 292 } 293 PetscScalar *XX = ptr; 294 ptr += stride; 295 if (idx++ == nShareVec) { 296 ptr = work_space_global; 297 stride = stride_global; 298 } 299 PetscScalar *R = ptr; 300 ptr += stride; 301 if (idx++ == nShareVec) { 302 ptr = work_space_global; 303 stride = stride_global; 304 } 305 PetscScalar *RP = ptr; 306 ptr += stride; 307 if (idx++ == nShareVec) { 308 ptr = work_space_global; 309 stride = stride_global; 310 } 311 PetscScalar *V = ptr; 312 ptr += stride; 313 if (idx++ == nShareVec) { 314 ptr = work_space_global; 315 stride = stride_global; 316 } 317 PetscScalar *T = ptr; 318 ptr += stride; 319 if (idx++ == nShareVec) { 320 ptr = work_space_global; 321 stride = stride_global; 322 } 323 PetscScalar *Q = ptr; 324 ptr += stride; 325 if (idx++ == nShareVec) { 326 ptr = work_space_global; 327 stride = stride_global; 328 } 329 PetscScalar *P = ptr; 330 ptr += stride; 331 if (idx++ == nShareVec) { 332 ptr = work_space_global; 333 stride = stride_global; 334 } 335 PetscScalar *U = ptr; 336 ptr += stride; 337 if (idx++ == nShareVec) { 338 ptr = work_space_global; 339 stride = stride_global; 340 } 341 PetscScalar *D = ptr; 342 ptr += stride; 343 if (idx++ == nShareVec) { 344 ptr = work_space_global; 345 stride = stride_global; 346 } 347 PetscScalar *AUQ = V; 348 349 // init: get b, zero x 350 parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) { 351 int rowa = ic[rowb]; 352 R[rowb - start] = glb_b[rowa]; 353 XX[rowb - start] = 0; 354 }); 355 team.team_barrier(); 356 parallel_reduce( 357 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += R[idx] * PetscConj(R[idx]); }, dpi); 358 team.team_barrier(); 359 r0 = dp = PetscSqrtReal(PetscRealPart(dpi)); 360 // diagnostics 361 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 362 if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e \n", 0, (double)dp); }); 363 #endif 364 if (dp < atol) { 365 metad->reason = KSP_CONVERGED_ATOL_NORMAL; 366 return 0; 367 } 368 if (0 == maxit) { 369 metad->reason = KSP_DIVERGED_ITS; 370 return 0; 371 } 372 373 /* Make the initial Rp = R */ 374 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { RP[idx] = R[idx]; }); 375 team.team_barrier(); 376 /* Set the initial conditions */ 377 etaold = 0.0; 378 psiold = 0.0; 379 tau = dp; 380 dpold = dp; 381 382 /* rhoold = (r,rp) */ 383 parallel_reduce( 384 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += R[idx] * PetscConj(RP[idx]); }, rhoold); 385 team.team_barrier(); 386 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { 387 U[idx] = R[idx]; 388 P[idx] = R[idx]; 389 T[idx] = Diag[idx] * P[idx]; 390 D[idx] = 0; 391 }); 392 team.team_barrier(); 393 MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, V); 394 395 i = 0; 396 do { 397 /* s <- (v,rp) */ 398 parallel_reduce( 399 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += V[idx] * PetscConj(RP[idx]); }, s); 400 team.team_barrier(); 401 a = rhoold / s; /* a <- rho / s */ 402 /* q <- u - a v VecWAXPY(w,alpha,x,y): w = alpha x + y. */ 403 /* t <- u + q */ 404 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { 405 Q[idx] = U[idx] - a * V[idx]; 406 T[idx] = U[idx] + Q[idx]; 407 }); 408 team.team_barrier(); 409 // KSP_PCApplyBAorAB 410 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { T[idx] = Diag[idx] * T[idx]; }); 411 team.team_barrier(); 412 MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, AUQ); 413 /* r <- r - a K (u + q) */ 414 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { R[idx] = R[idx] - a * AUQ[idx]; }); 415 team.team_barrier(); 416 parallel_reduce( 417 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += R[idx] * PetscConj(R[idx]); }, dpi); 418 team.team_barrier(); 419 dp = PetscSqrtReal(PetscRealPart(dpi)); 420 for (m = 0; m < 2; m++) { 421 if (!m) w = PetscSqrtReal(dp * dpold); 422 else w = dp; 423 psi = w / tau; 424 cm = 1.0 / PetscSqrtReal(1.0 + psi * psi); 425 tau = tau * psi * cm; 426 eta = cm * cm * a; 427 cf = psiold * psiold * etaold / a; 428 if (!m) { 429 /* D = U + cf D */ 430 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { D[idx] = U[idx] + cf * D[idx]; }); 431 } else { 432 /* D = Q + cf D */ 433 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { D[idx] = Q[idx] + cf * D[idx]; }); 434 } 435 team.team_barrier(); 436 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { XX[idx] = XX[idx] + eta * D[idx]; }); 437 team.team_barrier(); 438 dpest = PetscSqrtReal(2 * i + m + 2.0) * tau; 439 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 440 if (monitor && m == 1) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e \n", i + 1, (double)dpest); }); 441 #endif 442 if (dpest < atol) { 443 metad->reason = KSP_CONVERGED_ATOL_NORMAL; 444 goto done; 445 } 446 if (dpest / r0 < rtol) { 447 metad->reason = KSP_CONVERGED_RTOL_NORMAL; 448 goto done; 449 } 450 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 451 if (dpest / r0 > dtol) { 452 metad->reason = KSP_DIVERGED_DTOL; 453 Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: %d it, res=%e, r_0=%e\n", team.league_rank(), i, dpest, r0); }); 454 goto done; 455 } 456 #else 457 if (dpest / r0 > dtol) { 458 metad->reason = KSP_DIVERGED_DTOL; 459 goto done; 460 } 461 #endif 462 if (i + 1 == maxit) { 463 metad->reason = KSP_DIVERGED_ITS; 464 goto done; 465 } 466 467 etaold = eta; 468 psiold = psi; 469 } 470 471 /* rho <- (r,rp) */ 472 parallel_reduce( 473 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += R[idx] * PetscConj(RP[idx]); }, rho); 474 team.team_barrier(); 475 b = rho / rhoold; /* b <- rho / rhoold */ 476 /* u <- r + b q */ 477 /* p <- u + b(q + b p) */ 478 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { 479 U[idx] = R[idx] + b * Q[idx]; 480 Q[idx] = Q[idx] + b * P[idx]; 481 P[idx] = U[idx] + b * Q[idx]; 482 }); 483 /* v <- K p */ 484 team.team_barrier(); 485 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { T[idx] = Diag[idx] * P[idx]; }); 486 team.team_barrier(); 487 MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, V); 488 489 rhoold = rho; 490 dpold = dp; 491 492 i++; 493 } while (i < maxit); 494 done: 495 // KSPUnwindPreconditioner 496 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { XX[idx] = Diag[idx] * XX[idx]; }); 497 team.team_barrier(); 498 // put x into Plex order 499 parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) { 500 int rowa = ic[rowb]; 501 glb_x[rowa] = XX[rowb - start]; 502 }); 503 metad->its = i + 1; 504 if (1) { 505 int nnz; 506 parallel_reduce( 507 Kokkos::TeamVectorRange(team, start, end), [=](const int idx, int &lsum) { lsum += (glb_Aai[idx + 1] - glb_Aai[idx]); }, nnz); 508 metad->flops = 2 * (metad->its * (10 * Nblk + 2 * nnz) + 5 * Nblk); 509 } else { 510 metad->flops = 2 * (metad->its * (10 * Nblk + 2 * 50 * Nblk) + 5 * Nblk); // guess 511 } 512 return 0; 513 } 514 515 // Solve Ax = y with biCG 516 KOKKOS_INLINE_FUNCTION PetscErrorCode BJSolve_BICG(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, PetscScalar *work_space_global, const int stride_global, const int nShareVec, PetscScalar *work_space_shared, const int stride_shared, PetscReal rtol, PetscReal atol, PetscReal dtol, PetscInt maxit, Batch_MetaData *metad, const PetscInt start, const PetscInt end, const PetscScalar glb_idiag[], const PetscScalar *glb_b, PetscScalar *glb_x, bool monitor) { 517 using Kokkos::parallel_for; 518 using Kokkos::parallel_reduce; 519 int Nblk = end - start, i, stride = stride_shared, idx = 0; // start in shared mem 520 PetscReal dp, r0; 521 const PetscScalar *Di = &glb_idiag[start]; 522 PetscScalar *ptr = work_space_shared, dpi, a = 1.0, beta, betaold = 1.0, b, b2, ma, mac; 523 524 if (idx++ == nShareVec) { 525 ptr = work_space_global; 526 stride = stride_global; 527 } 528 PetscScalar *XX = ptr; 529 ptr += stride; 530 if (idx++ == nShareVec) { 531 ptr = work_space_global; 532 stride = stride_global; 533 } 534 PetscScalar *Rl = ptr; 535 ptr += stride; 536 if (idx++ == nShareVec) { 537 ptr = work_space_global; 538 stride = stride_global; 539 } 540 PetscScalar *Zl = ptr; 541 ptr += stride; 542 if (idx++ == nShareVec) { 543 ptr = work_space_global; 544 stride = stride_global; 545 } 546 PetscScalar *Pl = ptr; 547 ptr += stride; 548 if (idx++ == nShareVec) { 549 ptr = work_space_global; 550 stride = stride_global; 551 } 552 PetscScalar *Rr = ptr; 553 ptr += stride; 554 if (idx++ == nShareVec) { 555 ptr = work_space_global; 556 stride = stride_global; 557 } 558 PetscScalar *Zr = ptr; 559 ptr += stride; 560 if (idx++ == nShareVec) { 561 ptr = work_space_global; 562 stride = stride_global; 563 } 564 PetscScalar *Pr = ptr; 565 ptr += stride; 566 567 /* r <- b (x is 0) */ 568 parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) { 569 int rowa = ic[rowb]; 570 //PetscCall(VecCopy(Rr,Rl)); 571 Rl[rowb - start] = Rr[rowb - start] = glb_b[rowa]; 572 XX[rowb - start] = 0; 573 }); 574 team.team_barrier(); 575 /* z <- Br */ 576 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { 577 Zr[idx] = Di[idx] * Rr[idx]; 578 Zl[idx] = Di[idx] * Rl[idx]; 579 }); 580 team.team_barrier(); 581 /* dp <- r'*r */ 582 parallel_reduce( 583 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Rr[idx] * PetscConj(Rr[idx]); }, dpi); 584 team.team_barrier(); 585 r0 = dp = PetscSqrtReal(PetscRealPart(dpi)); 586 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 587 if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e \n", 0, (double)dp); }); 588 #endif 589 if (dp < atol) { 590 metad->reason = KSP_CONVERGED_ATOL_NORMAL; 591 return 0; 592 } 593 if (0 == maxit) { 594 metad->reason = KSP_DIVERGED_ITS; 595 return 0; 596 } 597 i = 0; 598 do { 599 /* beta <- r'z */ 600 parallel_reduce( 601 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += Zr[idx] * PetscConj(Rl[idx]); }, beta); 602 team.team_barrier(); 603 #if PCBJKOKKOS_VERBOSE_LEVEL >= 6 604 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 605 Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%7d beta = Z.R = %22.14e \n", i, (double)beta); }); 606 #endif 607 #endif 608 if (!i) { 609 if (beta == 0.0) { 610 metad->reason = KSP_DIVERGED_BREAKDOWN_BICG; 611 goto done; 612 } 613 /* p <- z */ 614 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { 615 Pr[idx] = Zr[idx]; 616 Pl[idx] = Zl[idx]; 617 }); 618 } else { 619 b = beta / betaold; 620 /* p <- z + b* p */ 621 b2 = PetscConj(b); 622 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { 623 Pr[idx] = b * Pr[idx] + Zr[idx]; 624 Pl[idx] = b2 * Pl[idx] + Zl[idx]; 625 }); 626 } 627 team.team_barrier(); 628 betaold = beta; 629 /* z <- Kp */ 630 MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, Pr, Zr); 631 MatMultTranspose(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, Pl, Zl); 632 /* dpi <- z'p */ 633 parallel_reduce( 634 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Zr[idx] * PetscConj(Pl[idx]); }, dpi); 635 team.team_barrier(); 636 // 637 a = beta / dpi; /* a = beta/p'z */ 638 ma = -a; 639 mac = PetscConj(ma); 640 /* x <- x + ap */ 641 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { 642 XX[idx] = XX[idx] + a * Pr[idx]; 643 Rr[idx] = Rr[idx] + ma * Zr[idx]; 644 Rl[idx] = Rl[idx] + mac * Zl[idx]; 645 }); 646 team.team_barrier(); 647 team.team_barrier(); 648 /* dp <- r'*r */ 649 parallel_reduce( 650 Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Rr[idx] * PetscConj(Rr[idx]); }, dpi); 651 team.team_barrier(); 652 dp = PetscSqrtReal(PetscRealPart(dpi)); 653 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 654 if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e \n", i + 1, (double)dp); }); 655 #endif 656 if (dp < atol) { 657 metad->reason = KSP_CONVERGED_ATOL_NORMAL; 658 goto done; 659 } 660 if (dp / r0 < rtol) { 661 metad->reason = KSP_CONVERGED_RTOL_NORMAL; 662 goto done; 663 } 664 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 665 if (dp / r0 > dtol) { 666 metad->reason = KSP_DIVERGED_DTOL; 667 Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: %d it, res=%e, r_0=%e\n", team.league_rank(), i, dp, r0); }); 668 goto done; 669 } 670 #else 671 if (dp / r0 > dtol) { 672 metad->reason = KSP_DIVERGED_DTOL; 673 goto done; 674 } 675 #endif 676 if (i + 1 == maxit) { 677 metad->reason = KSP_DIVERGED_ITS; 678 goto done; 679 } 680 /* z <- Br */ 681 parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { 682 Zr[idx] = Di[idx] * Rr[idx]; 683 Zl[idx] = Di[idx] * Rl[idx]; 684 }); 685 i++; 686 team.team_barrier(); 687 } while (i < maxit); 688 done: 689 // put x back into Plex order 690 parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) { 691 int rowa = ic[rowb]; 692 glb_x[rowa] = XX[rowb - start]; 693 }); 694 metad->its = i + 1; 695 if (1) { 696 int nnz; 697 parallel_reduce( 698 Kokkos::TeamVectorRange(team, start, end), [=](const int idx, int &lsum) { lsum += (glb_Aai[idx + 1] - glb_Aai[idx]); }, nnz); 699 metad->flops = 2 * (metad->its * (10 * Nblk + 2 * nnz) + 5 * Nblk); 700 } else { 701 metad->flops = 2 * (metad->its * (10 * Nblk + 2 * 50 * Nblk) + 5 * Nblk); // guess 702 } 703 return 0; 704 } 705 706 // KSP solver solve Ax = b; x is output, bin is input 707 static PetscErrorCode PCApply_BJKOKKOS(PC pc, Vec bin, Vec xout) { 708 PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data; 709 Mat A = pc->pmat; 710 Mat_SeqAIJKokkos *aijkok; 711 712 PetscFunctionBegin; 713 PetscCheck(jac->vec_diag && A, PetscObjectComm((PetscObject)pc), PETSC_ERR_USER, "Not setup???? %p %p", jac->vec_diag, A); 714 aijkok = static_cast<Mat_SeqAIJKokkos *>(A->spptr); 715 if (!aijkok) { 716 SETERRQ(PetscObjectComm((PetscObject)pc), PETSC_ERR_USER, "No aijkok"); 717 } else { 718 PetscInt maxit = jac->ksp->max_it; 719 const PetscInt conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp == 0 && PCBJKOKKOS_VEC_SIZE != 1) ? PCBJKOKKOS_TEAM_SIZE : 1; 720 const PetscInt nwork = jac->nwork, nBlk = jac->nBlocks; 721 PetscScalar *glb_xdata = NULL; 722 PetscReal rtol = jac->ksp->rtol, atol = jac->ksp->abstol, dtol = jac->ksp->divtol; 723 const PetscScalar *glb_idiag = jac->d_idiag_k->data(), *glb_bdata = NULL; 724 const PetscInt *glb_Aai = aijkok->i_device_data(), *glb_Aaj = aijkok->j_device_data(), *d_bid_eqOffset = jac->d_bid_eqOffset_k->data(); 725 const PetscScalar *glb_Aaa = aijkok->a_device_data(); 726 const PetscInt *d_isicol = jac->d_isicol_k->data(), *d_isrow = jac->d_isrow_k->data(); 727 PCFailedReason pcreason; 728 KSPIndex ksp_type_idx = jac->ksp_type_idx; 729 PetscMemType mtype; 730 PetscContainer container; 731 PetscInt batch_sz; 732 VecScatter plex_batch = NULL; // not used 733 Vec bvec; // a copy of b for scatter (just alias to bin now) 734 PetscBool monitor = jac->monitor; // captured 735 PetscInt view_bid = jac->batch_target; 736 MatInfo info; 737 jac->max_nits = 0; 738 if (view_bid < 0) view_bid = 0; 739 PetscCall(MatGetInfo(A, MAT_LOCAL, &info)); 740 // get field major is to map plex IO to/from block/field major 741 PetscCall(PetscObjectQuery((PetscObject)A, "plex_batch_is", (PetscObject *)&container)); 742 if (container) { 743 PetscCall(VecDuplicate(bin, &bvec)); 744 PetscCall(PetscContainerGetPointer(container, (void **)&plex_batch)); 745 PetscCall(VecScatterBegin(plex_batch, bin, bvec, INSERT_VALUES, SCATTER_FORWARD)); 746 PetscCall(VecScatterEnd(plex_batch, bin, bvec, INSERT_VALUES, SCATTER_FORWARD)); 747 SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "No plex_batch_is -- require NO field major ordering for now"); 748 } else { 749 bvec = bin; 750 } 751 // get x 752 PetscCall(VecGetArrayAndMemType(xout, &glb_xdata, &mtype)); 753 #if defined(PETSC_HAVE_CUDA) 754 PetscCheck(PetscMemTypeDevice(mtype), PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "No GPU data for x %" PetscInt_FMT " != %" PetscInt_FMT, mtype, PETSC_MEMTYPE_DEVICE); 755 #endif 756 PetscCall(VecGetArrayReadAndMemType(bvec, &glb_bdata, &mtype)); 757 #if defined(PETSC_HAVE_CUDA) 758 PetscCheck(PetscMemTypeDevice(mtype), PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "No GPU data for b"); 759 #endif 760 // get batch size 761 PetscCall(PetscObjectQuery((PetscObject)A, "batch size", (PetscObject *)&container)); 762 if (container) { 763 PetscInt *pNf = NULL; 764 PetscCall(PetscContainerGetPointer(container, (void **)&pNf)); 765 batch_sz = *pNf; 766 } else batch_sz = 1; 767 PetscCheck(nBlk % batch_sz == 0, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "batch_sz = %" PetscInt_FMT ", nBlk = %" PetscInt_FMT, batch_sz, nBlk); 768 if (ksp_type_idx == BATCH_KSP_GMRES_IDX) { // KK solver - move PETSc data into Kokkos Views, setup solver, solve, move data out of Kokkos, process metadata (convergence tests, etc.) 769 #if defined(PETSC_HAVE_KOKKOS_KERNELS_GMRES) 770 int Nsolves_team = jac->nsolves_team, fill_idx = 0; 771 int Nloc = jac->const_block_size; // same grids 772 const int Nsolves = nBlk; 773 const int nnz = (int)info.nz_used / Nsolves; // fix for variable grid size 774 if (Nsolves_team > batch_sz) Nsolves_team = batch_sz; // silently fix this 775 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)"); 776 PetscCheck(Nsolves % Nsolves_team == 0, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Nsolves.mod(Nsolves_team) != 0: Nsolves = %d, Nsolves_team = %d", Nsolves, Nsolves_team); 777 PetscCheck(((int)info.nz_used) % Nsolves == 0, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "info.nz_used.mod(Nsolves) != 0: info.nz_used = %g, Nsolves = %d", info.nz_used, Nsolves); 778 #if defined(PETSC_HAVE_CUDA_NVTX) 779 nvtxRangePushA("gmres-kk"); 780 #endif 781 Kokkos::View<PetscScalar **, layout, exec_space, Kokkos::MemoryTraits<Kokkos::Unmanaged>> inv_diag((PetscScalar *)glb_idiag, Nsolves, Nloc); // in correct order 782 if (!jac->rowOffsets) { 783 jac->rowOffsets = new IntView("rowOffsets", Nsolves / Nsolves_team, Nloc + 1); // same grids 784 jac->colIndices = new IntView("colIndices", Nsolves / Nsolves_team, nnz); 785 jac->batch_b = new XYType("batch rhs", Nsolves, Nloc); 786 jac->batch_x = new XYType("batch sol", Nsolves, Nloc); 787 jac->batch_values = new AMatrixValueView("batch values", Nsolves, nnz); 788 fill_idx = 1; 789 PetscInfo(pc, "Setup indices Nloc=%d, nnz=%d\n", Nloc, nnz); 790 } 791 IntView &rowOffsets = *jac->rowOffsets; 792 IntView &colIndices = *jac->colIndices; 793 XYType &batch_b = *jac->batch_b; 794 XYType &batch_x = *jac->batch_x; 795 AMatrixValueView &batch_values = *jac->batch_values; 796 797 Kokkos::deep_copy(batch_x, 0.); 798 PetscInfo(pc, "\tjac->n = %" PetscInt_FMT ", Nloc = %d, Nsolves = %d, nnz = %d, Nsolves_team = %d, league size = %d, maxit = %" PetscInt_FMT "\n", jac->n, Nloc, Nsolves, nnz, Nsolves_team, Nsolves / Nsolves_team, maxit); 799 Kokkos::parallel_for( 800 "rowOffsets+map", Kokkos::TeamPolicy<>(Nsolves, team_size, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) { 801 const int blkID = team.league_rank(), start = d_bid_eqOffset[blkID], end = d_bid_eqOffset[blkID + 1]; 802 if (fill_idx) { 803 if (blkID % Nsolves_team == 0) { // first matrix on this member 804 Kokkos::parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](const int rowb) { // Nloc 805 int rowa = d_isicol[rowb]; 806 int n = glb_Aai[rowa + 1] - glb_Aai[rowa]; 807 rowOffsets(blkID / Nsolves_team, rowb + 1 - start) = n; // save sizes 808 }); 809 } 810 } 811 // map b into field major space 812 Kokkos::parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) { 813 int rowa = d_isicol[rowb]; 814 batch_b(blkID, rowb - start) = glb_bdata[rowa]; 815 }); 816 }); 817 Kokkos::fence(); 818 if (fill_idx) { 819 Kokkos::parallel_for( 820 "prefix sum", Kokkos::TeamPolicy<>(Nsolves / Nsolves_team, 1, 1), KOKKOS_LAMBDA(const team_member team) { 821 const int graphID = team.league_rank(); 822 rowOffsets(graphID, 0) = 0; 823 for (size_t i = 0; i < Nloc; ++i) rowOffsets(graphID, i + 1) += rowOffsets(graphID, i); 824 }); 825 Kokkos::fence(); 826 } 827 Kokkos::parallel_for( 828 "copy matrix", Kokkos::TeamPolicy<>(Nsolves /* /batch_sz */, team_size, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) { 829 const int blkID = team.league_rank(), start = d_bid_eqOffset[blkID], end = d_bid_eqOffset[blkID + 1], graphID = blkID / Nsolves_team; 830 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, start, end), [=](const int rowb) { 831 int rowa = d_isicol[rowb]; // global index 832 int n = glb_Aai[rowa + 1] - glb_Aai[rowa]; 833 const PetscInt *aj = glb_Aaj + glb_Aai[rowa]; 834 const PetscScalar *aa = glb_Aaa + glb_Aai[rowa]; 835 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, n), [=](const int &i) { 836 PetscScalar val = aa[i]; 837 if (fill_idx && blkID % Nsolves_team == 0) colIndices(graphID, rowOffsets(graphID, rowb - start) + i) = d_isrow[aj[i]] - blkID * Nloc; // local" global - block start 838 batch_values(blkID, rowOffsets(graphID, rowb - start) + i) = val; 839 }); 840 }); 841 }); 842 Kokkos::fence(); 843 // setup solver 844 using ScalarType = typename AMatrixValueView::non_const_value_type; 845 using MagnitudeType = typename Kokkos::Details::ArithTraits<ScalarType>::mag_type; 846 using NormViewType = Kokkos::View<MagnitudeType *, layout, exec_space>; 847 using Norm2DViewType = Kokkos::View<MagnitudeType **, layout, exec_space>; 848 using Scalar3DViewType = Kokkos::View<ScalarType ***, layout, exec_space>; 849 using IntViewType = Kokkos::View<int *, layout, exec_space>; 850 using KrylovHandleType = KokkosBatched::KrylovHandle<Norm2DViewType, IntViewType, Scalar3DViewType>; 851 const int n_iterations = maxit; 852 const int team_size = -1; 853 const int vector_length = -1; 854 const double tol = rtol; 855 const int ortho_strategy = 0; 856 KrylovHandleType handle(Nsolves, Nsolves_team, n_iterations, true); 857 handle.Arnoldi_view = Scalar3DViewType("", Nsolves, n_iterations, Nloc + n_iterations + 3); 858 // solve 859 double time = Functor_TestBatchedTeamVectorGMRES<exec_space, AMatrixValueView, IntView, XYType, KrylovHandleType>(batch_values, inv_diag, rowOffsets, colIndices, batch_x, batch_b, Nsolves_team, team_size, vector_length, n_iterations, tol, ortho_strategy, 0, handle) 860 .run(pc); 861 Kokkos::fence(); 862 // get data back 863 Kokkos::parallel_for( 864 "map", Kokkos::TeamPolicy<>(Nsolves /* /batch_sz */, team_size, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) { 865 const int blkID = team.league_rank(), start = d_bid_eqOffset[blkID], end = d_bid_eqOffset[blkID + 1]; // 0 866 // map x into Plex/PETSc 867 Kokkos::parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) { 868 int rowa = d_isicol[rowb]; 869 glb_xdata[rowa] = batch_x(blkID, rowb - start); 870 }); 871 }); 872 // output assume species major - clone from Kokkos solvers 873 #if PCBJKOKKOS_VERBOSE_LEVEL >= 3 874 #if PCBJKOKKOS_VERBOSE_LEVEL >= 4 875 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "Iterations\n")); 876 #else 877 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "max iterations per species (gmres) :")); 878 #endif 879 for (PetscInt dmIdx = 0, s = 0, head = 0; dmIdx < jac->num_dms; dmIdx += batch_sz) { 880 for (PetscInt f = 0, idx = head; f < jac->dm_Nf[dmIdx]; f++, s++, idx++) { 881 #if PCBJKOKKOS_VERBOSE_LEVEL >= 4 882 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%2D:", s)); 883 for (int bid = 0; bid < batch_sz; bid++) { PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%3D ", handle.get_iteration_host(idx + bid * jac->dm_Nf[dmIdx]))); } 884 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\n")); 885 #else 886 int count = 0, ii; 887 for (int bid = 0; bid < batch_sz; bid++) { 888 if ((ii = handle.get_iteration_host(idx + bid * jac->dm_Nf[dmIdx])) > count) count = ii; 889 } 890 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%3d", count)); 891 #endif 892 } 893 head += batch_sz * jac->dm_Nf[dmIdx]; 894 } 895 #if PCBJKOKKOS_VERBOSE_LEVEL == 3 896 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\n")); 897 #endif 898 #endif 899 // return error code, get max it 900 PetscInt count = 0, mbid = 0; 901 if (handle.is_converged_host()) { 902 pcreason = PC_NOERROR; 903 if (!jac->max_nits) { 904 for (int blkID = 0; blkID < nBlk; blkID++) { 905 if (handle.get_iteration_host(blkID) > jac->max_nits) { 906 jac->max_nits = handle.get_iteration_host(blkID); 907 mbid = blkID; 908 } 909 } 910 } 911 } else { 912 PetscCall(PetscPrintf(PETSC_COMM_SELF, "There is at least one system that did not converge.")); 913 pcreason = PC_SUBPC_ERROR; 914 } 915 // output - assume species major order 916 for (int blkID = 0; blkID < nBlk; blkID++) { 917 if (jac->reason) { // -pc_bjkokkos_ksp_converged_reason 918 if (jac->batch_target == blkID) { 919 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), " Linear solve %s in %d iterations, batch %" PetscInt_FMT ", species %" PetscInt_FMT "\n", handle.is_converged_host(blkID) ? "converged" : "diverged", handle.get_iteration_host(blkID), blkID % batch_sz, blkID / batch_sz)); 920 } else if (jac->batch_target == -1 && handle.get_iteration_host(blkID) > count) { 921 jac->max_nits = count = handle.get_iteration_host(blkID); 922 mbid = blkID; 923 } 924 if (!handle.is_converged_host(blkID)) { PetscCall(PetscPrintf(PETSC_COMM_SELF, "ERROR species %d, batch %d did not converge with %d iterations\n", (int)(blkID / batch_sz), (int)blkID % batch_sz, handle.get_iteration_host(blkID))); } 925 } 926 } 927 if (jac->batch_target == -1 && jac->reason) { 928 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), " Linear solve %s in %d iteration, batch %" PetscInt_FMT ", specie %" PetscInt_FMT "\n", handle.is_converged_host(mbid) ? "converged" : "diverged", jac->max_nits, mbid % batch_sz, mbid / batch_sz)); 929 } 930 #if defined(PETSC_HAVE_CUDA_NVTX) 931 nvtxRangePop(); 932 #endif 933 #else 934 SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "batch GMRES not supported"); 935 #endif 936 } else { // Kokkos Krylov 937 using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space; 938 using vect2D_scr_t = Kokkos::View<PetscScalar **, Kokkos::LayoutLeft, scr_mem_t>; 939 Kokkos::View<Batch_MetaData *, Kokkos::DefaultExecutionSpace> d_metadata("solver meta data", nBlk); 940 int stride_shared, stride_global, global_buff_words; 941 d_bid_eqOffset = jac->d_bid_eqOffset_k->data(); 942 // solve each block independently 943 int scr_bytes_team_shared = 0, nShareVec = 0, nGlobBVec = 0; 944 if (jac->const_block_size) { // use shared memory for work vectors only if constant block size - todo: test efficiency loss 945 int maximum_shared_mem_size; 946 stride_shared = jac->const_block_size; // captured 947 #if defined(PETSC_HAVE_CUDA) 948 int device; 949 cudaError_t ier = cudaGetDevice(&device); 950 ier = cudaDeviceGetAttribute(&maximum_shared_mem_size, cudaDevAttrMaxSharedMemoryPerBlock, device); 951 #elif defined(PETSC_HAVE_HIP) 952 int device; 953 hipGetDevice(&device); 954 hipDeviceGetAttribute(&maximum_shared_mem_size, hipDeviceAttributeMaxSharedMemoryPerBlock, device); 955 #elif defined(PETSC_HAVE_SYCL) 956 maximum_shared_mem_size = 64000; 957 #else 958 maximum_shared_mem_size = 72000; 959 #endif 960 nShareVec = maximum_shared_mem_size / (jac->const_block_size * sizeof(PetscScalar)); // integer floor, number of vectors that fit in shared 961 if (nShareVec > nwork) nShareVec = nwork; 962 else nGlobBVec = nwork - nShareVec; 963 global_buff_words = jac->n * nGlobBVec; 964 scr_bytes_team_shared = jac->const_block_size * nShareVec * sizeof(PetscScalar); 965 //PetscCall(PetscPrintf(PETSC_COMM_WORLD,"maximum_shared_mem_size=%d scr_bytes_shared=%d nShareVec=%d, nGlobBVec=%d vec size=%d jac->const_block_size=%d\n",maximum_shared_mem_size,scr_bytes_team_shared,nShareVec,nGlobBVec,jac->const_block_size*sizeof(PetscScalar),jac->const_block_size)); 966 } else { 967 scr_bytes_team_shared = 0; 968 stride_shared = 0; 969 global_buff_words = jac->n * nwork; 970 nGlobBVec = nwork; // not needed == fix 971 } 972 stride_global = jac->n; // captured 973 #if defined(PETSC_HAVE_CUDA_NVTX) 974 nvtxRangePushA("batch-kokkos-solve"); 975 #endif 976 Kokkos::View<PetscScalar *, Kokkos::DefaultExecutionSpace> d_work_vecs_k("workvectors", global_buff_words); // global work vectors 977 PetscInfo(pc, "\tn = %d. %d shared bytes/team, %d global mem bytes, rtol=%e, num blocks %d, team_size=%d, %d vector threads, %d shared vectors, %d global vectors\n", (int)jac->n, scr_bytes_team_shared, global_buff_words, rtol, (int)nBlk, (int)team_size, PCBJKOKKOS_VEC_SIZE, nShareVec, nGlobBVec); 978 PetscScalar *d_work_vecs = d_work_vecs_k.data(); 979 Kokkos::parallel_for( 980 "Solve", Kokkos::TeamPolicy<Kokkos::LaunchBounds<256, 4>>(nBlk, team_size, PCBJKOKKOS_VEC_SIZE).set_scratch_size(PCBJKOKKOS_SHARED_LEVEL, Kokkos::PerTeam(scr_bytes_team_shared)), KOKKOS_LAMBDA(const team_member team) { 981 const int blkID = team.league_rank(), start = d_bid_eqOffset[blkID], end = d_bid_eqOffset[blkID + 1]; 982 vect2D_scr_t work_vecs_shared(team.team_scratch(PCBJKOKKOS_SHARED_LEVEL), end - start, nShareVec); 983 PetscScalar *work_buff_shared = work_vecs_shared.data(); 984 PetscScalar *work_buff_global = &d_work_vecs[start]; // start inc'ed in 985 bool print = monitor && (blkID == view_bid); 986 switch (ksp_type_idx) { 987 case BATCH_KSP_BICG_IDX: 988 BJSolve_BICG(team, glb_Aai, glb_Aaj, glb_Aaa, d_isrow, d_isicol, work_buff_global, stride_global, nShareVec, work_buff_shared, stride_shared, rtol, atol, dtol, maxit, &d_metadata[blkID], start, end, glb_idiag, glb_bdata, glb_xdata, print); 989 break; 990 case BATCH_KSP_TFQMR_IDX: 991 BJSolve_TFQMR(team, glb_Aai, glb_Aaj, glb_Aaa, d_isrow, d_isicol, work_buff_global, stride_global, nShareVec, work_buff_shared, stride_shared, rtol, atol, dtol, maxit, &d_metadata[blkID], start, end, glb_idiag, glb_bdata, glb_xdata, print); 992 break; 993 case BATCH_KSP_GMRES_IDX: 994 //BJSolve_GMRES(); 995 break; 996 default: 997 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 998 printf("Unknown KSP type %d\n", ksp_type_idx); 999 #else 1000 /* void */; 1001 #endif 1002 } 1003 }); 1004 Kokkos::fence(); 1005 #if defined(PETSC_HAVE_CUDA_NVTX) 1006 nvtxRangePop(); 1007 nvtxRangePushA("Post-solve-metadata"); 1008 #endif 1009 auto h_metadata = Kokkos::create_mirror(Kokkos::HostSpace::memory_space(), d_metadata); 1010 Kokkos::deep_copy(h_metadata, d_metadata); 1011 #if PCBJKOKKOS_VERBOSE_LEVEL >= 3 1012 #if PCBJKOKKOS_VERBOSE_LEVEL >= 4 1013 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Iterations\n")); 1014 #endif 1015 // assume species major 1016 #if PCBJKOKKOS_VERBOSE_LEVEL < 4 1017 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "max iterations per species (%s) :", ksp_type_idx == BATCH_KSP_BICG_IDX ? "bicg" : "tfqmr")); 1018 #endif 1019 for (PetscInt dmIdx = 0, s = 0, head = 0; dmIdx < jac->num_dms; dmIdx += batch_sz) { 1020 for (PetscInt f = 0, idx = head; f < jac->dm_Nf[dmIdx]; f++, s++, idx++) { 1021 #if PCBJKOKKOS_VERBOSE_LEVEL >= 4 1022 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%2" PetscInt_FMT ":", s)); 1023 for (int bid = 0; bid < batch_sz; bid++) { PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%3" PetscInt_FMT " ", h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its)); } 1024 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\n")); 1025 #else 1026 PetscInt count = 0; 1027 for (int bid = 0; bid < batch_sz; bid++) { 1028 if (h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its > count) count = h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its; 1029 } 1030 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%3" PetscInt_FMT " ", count)); 1031 #endif 1032 } 1033 head += batch_sz * jac->dm_Nf[dmIdx]; 1034 } 1035 #if PCBJKOKKOS_VERBOSE_LEVEL == 3 1036 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\n")); 1037 #endif 1038 #endif 1039 PetscInt count = 0, mbid = 0; 1040 for (int blkID = 0; blkID < nBlk; blkID++) { 1041 PetscCall(PetscLogGpuFlops((PetscLogDouble)h_metadata[blkID].flops)); 1042 if (jac->reason) { // -pc_bjkokkos_ksp_converged_reason 1043 if (jac->batch_target == blkID) { 1044 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), " Linear solve converged due to %s iterations %d, batch %" PetscInt_FMT ", species %" PetscInt_FMT "\n", KSPConvergedReasons[h_metadata[blkID].reason], h_metadata[blkID].its, blkID % batch_sz, blkID / batch_sz)); 1045 } else if (jac->batch_target == -1 && h_metadata[blkID].its > count) { 1046 jac->max_nits = count = h_metadata[blkID].its; 1047 mbid = blkID; 1048 } 1049 if (h_metadata[blkID].reason < 0) { 1050 PetscCall(PetscPrintf(PETSC_COMM_SELF, "ERROR reason=%s, its=%" PetscInt_FMT ". species %" PetscInt_FMT ", batch %" PetscInt_FMT "\n", KSPConvergedReasons[h_metadata[blkID].reason], h_metadata[blkID].its, blkID / batch_sz, blkID % batch_sz)); 1051 } 1052 } 1053 } 1054 if (jac->batch_target == -1 && jac->reason) { 1055 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), " Linear solve converged due to %s iterations %d, batch %" PetscInt_FMT ", specie %" PetscInt_FMT "\n", KSPConvergedReasons[h_metadata[mbid].reason], h_metadata[mbid].its, mbid % batch_sz, mbid / batch_sz)); 1056 } 1057 { 1058 int errsum; 1059 Kokkos::parallel_reduce( 1060 nBlk, 1061 KOKKOS_LAMBDA(const int idx, int &lsum) { 1062 if (d_metadata[idx].reason < 0) ++lsum; 1063 }, 1064 errsum); 1065 pcreason = errsum ? PC_SUBPC_ERROR : PC_NOERROR; 1066 if (!errsum && !jac->max_nits) { // set max its to give back to top KSP 1067 for (int blkID = 0; blkID < nBlk; blkID++) { 1068 if (h_metadata[blkID].its > jac->max_nits) jac->max_nits = h_metadata[blkID].its; 1069 } 1070 } else if (errsum) { 1071 PetscCall(PetscPrintf(PETSC_COMM_SELF, "ERROR Kokkos batch solver did not converge in all solves\n")); 1072 } 1073 } 1074 #if defined(PETSC_HAVE_CUDA_NVTX) 1075 nvtxRangePop(); 1076 #endif 1077 } // end of Kokkos (not Kernels) solvers block 1078 PetscCall(VecRestoreArrayAndMemType(xout, &glb_xdata)); 1079 PetscCall(VecRestoreArrayReadAndMemType(bvec, &glb_bdata)); 1080 PetscCall(PCSetFailedReason(pc, pcreason)); 1081 // map back to Plex space - not used 1082 if (plex_batch) { 1083 PetscCall(VecCopy(xout, bvec)); 1084 PetscCall(VecScatterBegin(plex_batch, bvec, xout, INSERT_VALUES, SCATTER_REVERSE)); 1085 PetscCall(VecScatterEnd(plex_batch, bvec, xout, INSERT_VALUES, SCATTER_REVERSE)); 1086 PetscCall(VecDestroy(&bvec)); 1087 } 1088 } // whole 'have aijkok' block 1089 PetscFunctionReturn(0); 1090 } 1091 1092 static PetscErrorCode PCSetUp_BJKOKKOS(PC pc) { 1093 PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data; 1094 Mat A = pc->pmat; 1095 Mat_SeqAIJKokkos *aijkok; 1096 PetscBool flg; 1097 1098 PetscFunctionBegin; 1099 PetscCheck(!pc->useAmat, PetscObjectComm((PetscObject)pc), PETSC_ERR_SUP, "No support for using 'use_amat'"); 1100 PetscCheck(A, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "No matrix - A is used above"); 1101 PetscCall(PetscObjectTypeCompareAny((PetscObject)A, &flg, MATSEQAIJKOKKOS, MATMPIAIJKOKKOS, MATAIJKOKKOS, "")); 1102 PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "must use '-dm_mat_type aijkokkos -dm_vec_type kokkos' for -pc_type bjkokkos"); 1103 if (!(aijkok = static_cast<Mat_SeqAIJKokkos *>(A->spptr))) { 1104 SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "No aijkok"); 1105 } else { 1106 if (!jac->vec_diag) { 1107 Vec *subX; 1108 DM pack, *subDM; 1109 PetscInt nDMs, n; 1110 PetscContainer container; 1111 PetscCall(PetscObjectQuery((PetscObject)A, "plex_batch_is", (PetscObject *)&container)); 1112 { // Permute the matrix to get a block diagonal system: d_isrow_k, d_isicol_k 1113 MatOrderingType rtype; 1114 IS isrow, isicol; 1115 const PetscInt *rowindices, *icolindices; 1116 rtype = MATORDERINGRCM; 1117 // get permutation. Not what I expect so inverted here 1118 PetscCall(MatGetOrdering(A, rtype, &isrow, &isicol)); 1119 PetscCall(ISDestroy(&isrow)); 1120 PetscCall(ISInvertPermutation(isicol, PETSC_DECIDE, &isrow)); // THIS IS BACKWARD -- isrow is inverse -- FIX!!!!! 1121 1122 Mat mat_block_order; 1123 PetscCall(MatCreateSubMatrix(A, isicol, isicol, MAT_INITIAL_MATRIX, &mat_block_order)); 1124 PetscCall(MatViewFromOptions(mat_block_order, NULL, "-ksp_batch_reorder_view")); 1125 PetscCall(MatDestroy(&mat_block_order)); 1126 1127 PetscCall(ISGetIndices(isrow, &rowindices)); 1128 PetscCall(ISGetIndices(isicol, &icolindices)); 1129 const Kokkos::View<PetscInt *, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_isrow_k((PetscInt *)rowindices, A->rmap->n); 1130 const Kokkos::View<PetscInt *, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_isicol_k((PetscInt *)icolindices, A->rmap->n); 1131 jac->d_isrow_k = new Kokkos::View<PetscInt *>(Kokkos::create_mirror(DefaultMemorySpace(), h_isrow_k)); 1132 jac->d_isicol_k = new Kokkos::View<PetscInt *>(Kokkos::create_mirror(DefaultMemorySpace(), h_isicol_k)); 1133 Kokkos::deep_copy(*jac->d_isrow_k, h_isrow_k); 1134 Kokkos::deep_copy(*jac->d_isicol_k, h_isicol_k); 1135 PetscCall(ISRestoreIndices(isrow, &rowindices)); 1136 PetscCall(ISRestoreIndices(isicol, &icolindices)); 1137 PetscCall(ISDestroy(&isrow)); 1138 PetscCall(ISDestroy(&isicol)); 1139 } 1140 // get block sizes 1141 PetscCall(PCGetDM(pc, &pack)); 1142 PetscCheck(pack, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "no DM. Requires a composite DM"); 1143 PetscCall(PetscObjectTypeCompare((PetscObject)pack, DMCOMPOSITE, &flg)); 1144 PetscCheck(flg, PetscObjectComm((PetscObject)pack), PETSC_ERR_USER, "Not for type %s", ((PetscObject)pack)->type_name); 1145 PetscCall(DMCompositeGetNumberDM(pack, &nDMs)); 1146 jac->num_dms = nDMs; 1147 PetscCall(DMCreateGlobalVector(pack, &jac->vec_diag)); 1148 PetscCall(VecGetLocalSize(jac->vec_diag, &n)); 1149 jac->n = n; 1150 jac->d_idiag_k = new Kokkos::View<PetscScalar *, Kokkos::LayoutRight>("idiag", n); 1151 // options 1152 PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc)); 1153 PetscCall(KSPSetFromOptions(jac->ksp)); 1154 PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPBICG, "")); 1155 if (flg) { 1156 jac->ksp_type_idx = BATCH_KSP_BICG_IDX; 1157 jac->nwork = 7; 1158 } else { 1159 PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPTFQMR, "")); 1160 if (flg) { 1161 jac->ksp_type_idx = BATCH_KSP_TFQMR_IDX; 1162 jac->nwork = 10; 1163 } else { 1164 PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPGMRES, "")); 1165 if (flg) { 1166 jac->ksp_type_idx = BATCH_KSP_GMRES_IDX; 1167 jac->nwork = 0; 1168 } else { 1169 SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Unsupported batch ksp type"); 1170 } 1171 } 1172 } 1173 PetscOptionsBegin(PetscObjectComm((PetscObject)jac->ksp), ((PetscObject)jac->ksp)->prefix, "Options for Kokkos batch solver", "none"); 1174 PetscCall(PetscOptionsBool("-ksp_converged_reason", "", "bjkokkos.kokkos.cxx.c", jac->reason, &jac->reason, NULL)); 1175 PetscCall(PetscOptionsBool("-ksp_monitor", "", "bjkokkos.kokkos.cxx.c", jac->monitor, &jac->monitor, NULL)); 1176 PetscCall(PetscOptionsInt("-ksp_batch_target", "", "bjkokkos.kokkos.cxx.c", jac->batch_target, &jac->batch_target, NULL)); 1177 PetscCall(PetscOptionsInt("-ksp_batch_nsolves_team", "", "bjkokkos.kokkos.cxx.c", jac->nsolves_team, &jac->nsolves_team, NULL)); 1178 PetscCheck(jac->batch_target < jac->num_dms, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "-ksp_batch_target (%" PetscInt_FMT ") >= number of DMs (%" PetscInt_FMT ")", jac->batch_target, jac->num_dms); 1179 PetscOptionsEnd(); 1180 // get blocks - jac->d_bid_eqOffset_k 1181 PetscCall(PetscMalloc(sizeof(*subX) * nDMs, &subX)); 1182 PetscCall(PetscMalloc(sizeof(*subDM) * nDMs, &subDM)); 1183 PetscCall(PetscMalloc(sizeof(*jac->dm_Nf) * nDMs, &jac->dm_Nf)); 1184 PetscCall(PetscInfo(pc, "Have %" PetscInt_FMT " DMs, n=%" PetscInt_FMT " rtol=%g type = %s\n", nDMs, n, (double)jac->ksp->rtol, ((PetscObject)jac->ksp)->type_name)); 1185 PetscCall(DMCompositeGetEntriesArray(pack, subDM)); 1186 jac->nBlocks = 0; 1187 for (PetscInt ii = 0; ii < nDMs; ii++) { 1188 PetscSection section; 1189 PetscInt Nf; 1190 DM dm = subDM[ii]; 1191 PetscCall(DMGetLocalSection(dm, §ion)); 1192 PetscCall(PetscSectionGetNumFields(section, &Nf)); 1193 jac->nBlocks += Nf; 1194 #if PCBJKOKKOS_VERBOSE_LEVEL <= 2 1195 if (ii == 0) PetscCall(PetscInfo(pc, "%" PetscInt_FMT ") %" PetscInt_FMT " blocks (%" PetscInt_FMT " total)\n", ii, Nf, jac->nBlocks)); 1196 #else 1197 PetscCall(PetscInfo(pc, "%" PetscInt_FMT ") %" PetscInt_FMT " blocks (%" PetscInt_FMT " total)\n", ii, Nf, jac->nBlocks)); 1198 #endif 1199 jac->dm_Nf[ii] = Nf; 1200 } 1201 { // d_bid_eqOffset_k 1202 Kokkos::View<PetscInt *, Kokkos::LayoutRight, Kokkos::HostSpace> h_block_offsets("block_offsets", jac->nBlocks + 1); 1203 PetscCall(DMCompositeGetAccessArray(pack, jac->vec_diag, nDMs, NULL, subX)); 1204 h_block_offsets[0] = 0; 1205 jac->const_block_size = -1; 1206 for (PetscInt ii = 0, idx = 0; ii < nDMs; ii++) { 1207 PetscInt nloc, nblk; 1208 PetscCall(VecGetSize(subX[ii], &nloc)); 1209 nblk = nloc / jac->dm_Nf[ii]; 1210 PetscCheck(nloc % jac->dm_Nf[ii] == 0, PetscObjectComm((PetscObject)pc), PETSC_ERR_USER, "nloc%%jac->dm_Nf[ii] (%" PetscInt_FMT ") != 0 DMs", nloc % jac->dm_Nf[ii]); 1211 for (PetscInt jj = 0; jj < jac->dm_Nf[ii]; jj++, idx++) { 1212 h_block_offsets[idx + 1] = h_block_offsets[idx] + nblk; 1213 #if PCBJKOKKOS_VERBOSE_LEVEL <= 2 1214 if (idx == 0) PetscCall(PetscInfo(pc, "\t%" PetscInt_FMT ") Add block with %" PetscInt_FMT " equations of %" PetscInt_FMT "\n", idx + 1, nblk, jac->nBlocks)); 1215 #else 1216 PetscCall(PetscInfo(pc, "\t%" PetscInt_FMT ") Add block with %" PetscInt_FMT " equations of %" PetscInt_FMT "\n", idx + 1, nblk, jac->nBlocks)); 1217 #endif 1218 if (jac->const_block_size == -1) jac->const_block_size = nblk; 1219 else if (jac->const_block_size > 0 && jac->const_block_size != nblk) jac->const_block_size = 0; 1220 } 1221 } 1222 PetscCall(DMCompositeRestoreAccessArray(pack, jac->vec_diag, jac->nBlocks, NULL, subX)); 1223 PetscCall(PetscFree(subX)); 1224 PetscCall(PetscFree(subDM)); 1225 jac->d_bid_eqOffset_k = new Kokkos::View<PetscInt *, Kokkos::LayoutRight>(Kokkos::create_mirror(Kokkos::DefaultExecutionSpace::memory_space(), h_block_offsets)); 1226 Kokkos::deep_copy(*jac->d_bid_eqOffset_k, h_block_offsets); 1227 } 1228 } 1229 { // get jac->d_idiag_k (PC setup), 1230 const PetscInt *d_ai = aijkok->i_device_data(), *d_aj = aijkok->j_device_data(); 1231 const PetscScalar *d_aa = aijkok->a_device_data(); 1232 const PetscInt conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp == 0 && PCBJKOKKOS_VEC_SIZE != 1) ? PCBJKOKKOS_TEAM_SIZE : 1; 1233 PetscInt *d_bid_eqOffset = jac->d_bid_eqOffset_k->data(), *r = jac->d_isrow_k->data(), *ic = jac->d_isicol_k->data(); 1234 PetscScalar *d_idiag = jac->d_idiag_k->data(); 1235 Kokkos::parallel_for( 1236 "Diag", Kokkos::TeamPolicy<>(jac->nBlocks, team_size, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) { 1237 const PetscInt blkID = team.league_rank(); 1238 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, d_bid_eqOffset[blkID], d_bid_eqOffset[blkID + 1]), [=](const int rowb) { 1239 const PetscInt rowa = ic[rowb], ai = d_ai[rowa], *aj = d_aj + ai; // grab original data 1240 const PetscScalar *aa = d_aa + ai; 1241 const PetscInt nrow = d_ai[rowa + 1] - ai; 1242 int found; 1243 Kokkos::parallel_reduce( 1244 Kokkos::ThreadVectorRange(team, nrow), 1245 [=](const int &j, int &count) { 1246 const PetscInt colb = r[aj[j]]; 1247 if (colb == rowb) { 1248 d_idiag[rowb] = 1. / aa[j]; 1249 count++; 1250 } 1251 }, 1252 found); 1253 #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL) 1254 if (found != 1) Kokkos::single(Kokkos::PerThread(team), [=]() { printf("ERRORrow %d) found = %d\n", rowb, found); }); 1255 #endif 1256 }); 1257 }); 1258 } 1259 } 1260 PetscFunctionReturn(0); 1261 } 1262 1263 /* Default destroy, if it has never been setup */ 1264 static PetscErrorCode PCReset_BJKOKKOS(PC pc) { 1265 PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data; 1266 1267 PetscFunctionBegin; 1268 PetscCall(KSPDestroy(&jac->ksp)); 1269 PetscCall(VecDestroy(&jac->vec_diag)); 1270 if (jac->d_bid_eqOffset_k) delete jac->d_bid_eqOffset_k; 1271 if (jac->d_idiag_k) delete jac->d_idiag_k; 1272 if (jac->d_isrow_k) delete jac->d_isrow_k; 1273 if (jac->d_isicol_k) delete jac->d_isicol_k; 1274 jac->d_bid_eqOffset_k = NULL; 1275 jac->d_idiag_k = NULL; 1276 jac->d_isrow_k = NULL; 1277 jac->d_isicol_k = NULL; 1278 PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSGetKSP_C", NULL)); // not published now (causes configure errors) 1279 PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSSetKSP_C", NULL)); 1280 PetscCall(PetscFree(jac->dm_Nf)); 1281 jac->dm_Nf = NULL; 1282 if (jac->rowOffsets) delete jac->rowOffsets; 1283 if (jac->colIndices) delete jac->colIndices; 1284 if (jac->batch_b) delete jac->batch_b; 1285 if (jac->batch_x) delete jac->batch_x; 1286 if (jac->batch_values) delete jac->batch_values; 1287 jac->rowOffsets = NULL; 1288 jac->colIndices = NULL; 1289 jac->batch_b = NULL; 1290 jac->batch_x = NULL; 1291 jac->batch_values = NULL; 1292 1293 PetscFunctionReturn(0); 1294 } 1295 1296 static PetscErrorCode PCDestroy_BJKOKKOS(PC pc) { 1297 PetscFunctionBegin; 1298 PetscCall(PCReset_BJKOKKOS(pc)); 1299 PetscCall(PetscFree(pc->data)); 1300 PetscFunctionReturn(0); 1301 } 1302 1303 static PetscErrorCode PCView_BJKOKKOS(PC pc, PetscViewer viewer) { 1304 PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data; 1305 PetscBool iascii; 1306 1307 PetscFunctionBegin; 1308 if (!jac->ksp) PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc)); 1309 PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii)); 1310 if (iascii) { 1311 PetscCall(PetscViewerASCIIPrintf(viewer, " Batched device linear solver: Krylov (KSP) method with Jacobi preconditioning\n")); 1312 PetscCall(PetscViewerASCIIPrintf(viewer, "\t\tnwork = %" PetscInt_FMT ", rel tol = %e, abs tol = %e, div tol = %e, max it =%" PetscInt_FMT ", type = %s\n", jac->nwork, jac->ksp->rtol, jac->ksp->abstol, jac->ksp->divtol, jac->ksp->max_it, 1313 ((PetscObject)jac->ksp)->type_name)); 1314 } 1315 PetscFunctionReturn(0); 1316 } 1317 1318 static PetscErrorCode PCSetFromOptions_BJKOKKOS(PC pc, PetscOptionItems *PetscOptionsObject) { 1319 PetscFunctionBegin; 1320 PetscOptionsHeadBegin(PetscOptionsObject, "PC BJKOKKOS options"); 1321 PetscOptionsHeadEnd(); 1322 PetscFunctionReturn(0); 1323 } 1324 1325 static PetscErrorCode PCBJKOKKOSSetKSP_BJKOKKOS(PC pc, KSP ksp) { 1326 PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data; 1327 1328 PetscFunctionBegin; 1329 PetscCall(PetscObjectReference((PetscObject)ksp)); 1330 PetscCall(KSPDestroy(&jac->ksp)); 1331 jac->ksp = ksp; 1332 PetscFunctionReturn(0); 1333 } 1334 1335 /*@C 1336 PCBJKOKKOSSetKSP - Sets the KSP context for a KSP PC. 1337 1338 Collective on PC 1339 1340 Input Parameters: 1341 + pc - the preconditioner context 1342 - ksp - the KSP solver 1343 1344 Notes: 1345 The PC and the KSP must have the same communicator 1346 1347 Level: advanced 1348 1349 @*/ 1350 PetscErrorCode PCBJKOKKOSSetKSP(PC pc, KSP ksp) { 1351 PetscFunctionBegin; 1352 PetscValidHeaderSpecific(pc, PC_CLASSID, 1); 1353 PetscValidHeaderSpecific(ksp, KSP_CLASSID, 2); 1354 PetscCheckSameComm(pc, 1, ksp, 2); 1355 PetscTryMethod(pc, "PCBJKOKKOSSetKSP_C", (PC, KSP), (pc, ksp)); 1356 PetscFunctionReturn(0); 1357 } 1358 1359 static PetscErrorCode PCBJKOKKOSGetKSP_BJKOKKOS(PC pc, KSP *ksp) { 1360 PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data; 1361 1362 PetscFunctionBegin; 1363 if (!jac->ksp) PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc)); 1364 *ksp = jac->ksp; 1365 PetscFunctionReturn(0); 1366 } 1367 1368 /*@C 1369 PCBJKOKKOSGetKSP - Gets the KSP context for a KSP PC. 1370 1371 Not Collective but KSP returned is parallel if PC was parallel 1372 1373 Input Parameter: 1374 . pc - the preconditioner context 1375 1376 Output Parameters: 1377 . ksp - the KSP solver 1378 1379 Notes: 1380 You must call KSPSetUp() before calling PCBJKOKKOSGetKSP(). 1381 1382 If the PC is not a PCBJKOKKOS object it raises an error 1383 1384 Level: advanced 1385 1386 @*/ 1387 PetscErrorCode PCBJKOKKOSGetKSP(PC pc, KSP *ksp) { 1388 PetscFunctionBegin; 1389 PetscValidHeaderSpecific(pc, PC_CLASSID, 1); 1390 PetscValidPointer(ksp, 2); 1391 PetscUseMethod(pc, "PCBJKOKKOSGetKSP_C", (PC, KSP *), (pc, ksp)); 1392 PetscFunctionReturn(0); 1393 } 1394 1395 static PetscErrorCode PCPostSolve_BJKOKKOS(PC pc, KSP ksp, Vec b, Vec x) { 1396 PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data; 1397 1398 PetscFunctionBegin; 1399 ksp->its = jac->max_nits; 1400 PetscFunctionReturn(0); 1401 } 1402 1403 /* ----------------------------------------------------------------------------------*/ 1404 1405 /*MC 1406 PCBJKOKKOS - Defines a preconditioner that applies a Krylov solver and preconditioner to the blocks in a AIJASeq matrix on the GPU. 1407 1408 Options Database Key: 1409 . -pc_bjkokkos_ - options prefix with ksp options 1410 1411 Level: intermediate 1412 1413 Notes: 1414 For use with -ksp_type preonly to bypass any CPU work 1415 1416 Developer Notes: 1417 1418 .seealso: `PCCreate()`, `PCSetType()`, `PCType`, `PC`, 1419 `PCSHELL`, `PCCOMPOSITE`, `PCSetUseAmat()`, `PCBJKOKKOSGetKSP()` 1420 1421 M*/ 1422 1423 PETSC_EXTERN PetscErrorCode PCCreate_BJKOKKOS(PC pc) { 1424 PC_PCBJKOKKOS *jac; 1425 1426 PetscFunctionBegin; 1427 PetscCall(PetscNewLog(pc, &jac)); 1428 pc->data = (void *)jac; 1429 1430 jac->ksp = NULL; 1431 jac->vec_diag = NULL; 1432 jac->d_bid_eqOffset_k = NULL; 1433 jac->d_idiag_k = NULL; 1434 jac->d_isrow_k = NULL; 1435 jac->d_isicol_k = NULL; 1436 jac->nBlocks = 1; 1437 jac->max_nits = 0; 1438 1439 PetscCall(PetscMemzero(pc->ops, sizeof(struct _PCOps))); 1440 pc->ops->apply = PCApply_BJKOKKOS; 1441 pc->ops->applytranspose = NULL; 1442 pc->ops->setup = PCSetUp_BJKOKKOS; 1443 pc->ops->reset = PCReset_BJKOKKOS; 1444 pc->ops->destroy = PCDestroy_BJKOKKOS; 1445 pc->ops->setfromoptions = PCSetFromOptions_BJKOKKOS; 1446 pc->ops->view = PCView_BJKOKKOS; 1447 pc->ops->postsolve = PCPostSolve_BJKOKKOS; 1448 1449 jac->rowOffsets = NULL; 1450 jac->colIndices = NULL; 1451 jac->batch_b = NULL; 1452 jac->batch_x = NULL; 1453 jac->batch_values = NULL; 1454 1455 PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSGetKSP_C", PCBJKOKKOSGetKSP_BJKOKKOS)); 1456 PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSSetKSP_C", PCBJKOKKOSSetKSP_BJKOKKOS)); 1457 PetscFunctionReturn(0); 1458 } 1459