Actual source code: pcbjkokkosimpl.h
1: #pragma once
3: #include <petscvec_kokkos.hpp>
4: #include <petsc/private/pcimpl.h>
5: #include <petsc/private/deviceimpl.h>
6: #include <petsc/private/kspimpl.h>
8: #include "Kokkos_Core.hpp"
10: #if defined(PETSC_HAVE_CUDA)
11: #include <nvToolsExt.h>
12: #endif
14: #define PCBJKOKKOS_SHARED_LEVEL 1 // 0 is shared, 1 is global
15: #define PCBJKOKKOS_VEC_SIZE 16
16: #define PCBJKOKKOS_TEAM_SIZE 16
18: #define PCBJKOKKOS_VERBOSE_LEVEL 1
20: typedef enum {
21: BATCH_KSP_BICG_IDX,
22: BATCH_KSP_TFQMR_IDX,
23: BATCH_KSP_GMRESKK_IDX,
24: BATCH_KSP_PREONLY_IDX,
25: NUM_BATCH_TYPES
26: } KSPIndex;
28: typedef Kokkos::DefaultExecutionSpace exec_space;
29: using layout = Kokkos::LayoutRight;
30: using IntView = Kokkos::View<PetscInt **, layout, exec_space>;
31: using AMatrixValueView = const Kokkos::View<PetscScalar **, layout, exec_space>;
32: using XYType = const Kokkos::View<PetscScalar **, layout, exec_space>;
34: typedef struct {
35: Vec vec_diag;
36: PetscInt nBlocks; /* total number of blocks */
37: PetscInt n; // cache host version of d_bid_eqOffset_k[nBlocks]
38: KSP ksp; // Used just for options. Should have one for each block
39: Kokkos::View<PetscInt *, Kokkos::LayoutRight> *d_bid_eqOffset_k;
40: Kokkos::View<PetscScalar *, Kokkos::LayoutRight> *d_idiag_k;
41: Kokkos::View<PetscInt *> *d_isrow_k;
42: Kokkos::View<PetscInt *> *d_isicol_k;
43: KSPIndex ksp_type_idx;
44: PetscInt nwork;
45: PetscInt const_block_size; // used to decide to use shared memory for work vectors
46: PetscInt *dm_Nf; // Number of fields in each DM
47: PetscInt num_dms;
48: // diagnostics
49: PetscBool reason;
50: PetscBool monitor;
51: PetscInt batch_target;
52: PetscInt rank_target;
53: PetscInt nsolves_team;
54: PetscInt max_nits;
55: // caches
56: IntView *rowOffsets;
57: IntView *colIndices;
58: XYType *batch_b;
59: XYType *batch_x;
60: AMatrixValueView *batch_values;
61: } PC_PCBJKOKKOS;
63: typedef Kokkos::TeamPolicy<>::member_type team_member;
64: #if defined(PETSC_HAVE_KOKKOS_KERNELS_BATCH)
65: PETSC_INTERN PetscErrorCode PCApply_BJKOKKOSKERNELS(PC, const PetscScalar *, PetscScalar *, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt, MatInfo, const PetscInt, PCFailedReason *);
66: #endif