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