1: #pragma once

  3: #include <petscvec_kokkos.hpp>
  4: #include <petscpkg_version.h>
  5: #include <petsc/private/pcimpl.h>
  6: #include <petsc/private/deviceimpl.h>
  7: #include <petsc/private/kspimpl.h>

  9: #include "Kokkos_Core.hpp"

 11: #if PetscDefined(HAVE_CUDA)
 12:   #if PETSC_PKG_CUDA_VERSION_GE(10, 0, 0)
 13:     #include <nvtx3/nvToolsExt.h>
 14:   #else
 15:     #include <nvToolsExt.h>
 16:   #endif
 17: #endif

 19: #define PCBJKOKKOS_SHARED_LEVEL 1 // 0 is shared, 1 is global
 20: #define PCBJKOKKOS_VEC_SIZE     16
 21: #define PCBJKOKKOS_TEAM_SIZE    16

 23: #define PCBJKOKKOS_VERBOSE_LEVEL 1

 25: typedef enum {
 26:   BATCH_KSP_BICG_IDX,
 27:   BATCH_KSP_TFQMR_IDX,
 28:   BATCH_KSP_GMRESKK_IDX,
 29:   BATCH_KSP_PREONLY_IDX,
 30:   NUM_BATCH_TYPES
 31: } KSPIndex;

 33: typedef Kokkos::DefaultExecutionSpace exec_space;
 34: using layout           = Kokkos::LayoutRight;
 35: using IntView          = Kokkos::View<PetscInt **, layout, exec_space>;
 36: using AMatrixValueView = const Kokkos::View<PetscScalar **, layout, exec_space>;
 37: using XYType           = const Kokkos::View<PetscScalar **, layout, exec_space>;

 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  rank_target;
 58:   PetscInt  nsolves_team;
 59:   PetscInt  max_nits;
 60:   // caches
 61:   IntView          *rowOffsets;
 62:   IntView          *colIndices;
 63:   XYType           *batch_b;
 64:   XYType           *batch_x;
 65:   AMatrixValueView *batch_values;
 66: } PC_PCBJKOKKOS;

 68: typedef Kokkos::TeamPolicy<>::member_type team_member;
 69: #if PetscDefined(HAVE_KOKKOS_KERNELS_BATCH)
 70: 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 *);
 71: #endif