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