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