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