1 #ifndef PETSCLANDAU_H 2 #define PETSCLANDAU_H 3 4 #include <petscdmplex.h> /*I "petscdmplex.h" I*/ 5 #include <petscts.h> 6 7 PETSC_EXTERN PetscErrorCode DMPlexLandauPrintNorms(Vec, PetscInt); 8 PETSC_EXTERN PetscErrorCode DMPlexLandauCreateVelocitySpace(MPI_Comm, PetscInt, const char[], Vec *, Mat *, DM *); 9 PETSC_EXTERN PetscErrorCode DMPlexLandauDestroyVelocitySpace(DM *); 10 PETSC_EXTERN PetscErrorCode DMPlexLandauAccess(DM, Vec, PetscErrorCode (*)(DM, Vec, PetscInt, PetscInt, PetscInt, void *), void *); 11 PETSC_EXTERN PetscErrorCode DMPlexLandauAddMaxwellians(DM, Vec, PetscReal, PetscReal[], PetscReal[], PetscInt, PetscInt, PetscInt, void *); 12 PETSC_EXTERN PetscErrorCode DMPlexLandauCreateMassMatrix(DM dm, Mat *Amat); 13 PETSC_EXTERN PetscErrorCode DMPlexLandauIFunction(TS, PetscReal, Vec, Vec, Vec, void *); 14 PETSC_EXTERN PetscErrorCode DMPlexLandauIJacobian(TS, PetscReal, Vec, Vec, PetscReal, Mat, Mat, void *); 15 16 typedef int LandauIdx; 17 18 /* the Fokker-Planck-Landau context */ 19 #if !defined(LANDAU_MAX_SPECIES) 20 #if defined(PETSC_USE_DMLANDAU_2D) 21 #define LANDAU_MAX_SPECIES 10 22 #define LANDAU_MAX_GRIDS 3 23 #else 24 #define LANDAU_MAX_SPECIES 10 25 #define LANDAU_MAX_GRIDS 3 26 #endif 27 #else 28 #define LANDAU_MAX_GRIDS 3 29 #endif 30 31 #if !defined(LANDAU_MAX_Q) 32 #if defined(LANDAU_MAX_NQND) 33 #error "LANDAU_MAX_NQND but not LANDAU_MAX_Q. Use -DLANDAU_MAX_Q=4 for Q3 elements" 34 #endif 35 #if defined(PETSC_USE_DMLANDAU_2D) 36 #define LANDAU_MAX_Q 6 37 #else 38 #define LANDAU_MAX_Q 4 39 #endif 40 #else 41 #undef LANDAU_MAX_NQND 42 #endif 43 44 #if defined(PETSC_USE_DMLANDAU_2D) 45 #define LANDAU_MAX_Q_FACE LANDAU_MAX_Q 46 #define LANDAU_MAX_NQND (LANDAU_MAX_Q * LANDAU_MAX_Q) 47 #define LANDAU_MAX_BATCH_SZ 1024 48 #define LANDAU_DIM 2 49 #else 50 #define LANDAU_MAX_Q_FACE (LANDAU_MAX_Q * LANDAU_MAX_Q) 51 #define LANDAU_MAX_NQND (LANDAU_MAX_Q * LANDAU_MAX_Q * LANDAU_MAX_Q) 52 #define LANDAU_MAX_BATCH_SZ 64 53 #define LANDAU_DIM 3 54 #endif 55 56 typedef enum { 57 LANDAU_KOKKOS, 58 LANDAU_CPU 59 } LandauDeviceType; 60 61 // static data - will be "device" data 62 typedef struct { 63 void *invJ; // nip*dim*dim 64 void *D; // nq*nb*dim 65 void *B; // nq*nb 66 void *alpha; // ns 67 void *beta; // ns 68 void *invMass; // ns 69 void *w; // nip 70 void *x; // nip 71 void *y; // nip 72 void *z; // nip 73 void *Eq_m; // ns - dynamic 74 void *f; // nip*Nf - dynamic (IP) 75 void *dfdx; // nip*Nf - dynamic (IP) 76 void *dfdy; // nip*Nf - dynamic (IP) 77 void *dfdz; // nip*Nf - dynamic (IP) 78 int dim_, ns_, nip_, nq_, nb_; 79 void *NCells; // remove and use elem_offset - TODO 80 void *species_offset; // for each grid, but same for all batched vertices 81 void *mat_offset; // for each grid, but same for all batched vertices 82 void *elem_offset; // for each grid, but same for all batched vertices 83 void *ip_offset; // for each grid, but same for all batched vertices 84 void *ipf_offset; // for each grid, but same for all batched vertices 85 void *ipfdf_data; // for each grid, but same for all batched vertices 86 void *maps; // for each grid, but same for all batched vertices 87 // COO 88 void *coo_elem_offsets; 89 void *coo_elem_point_offsets; 90 void *coo_elem_fullNb; 91 void *coo_vals; 92 void *lambdas; 93 LandauIdx coo_n_cellsTot; 94 LandauIdx coo_size; 95 LandauIdx coo_max_fullnb; 96 } LandauStaticData; 97 98 typedef enum { 99 LANDAU_EX2_TSSOLVE, 100 LANDAU_MATRIX_TOTAL, 101 LANDAU_OPERATOR, 102 LANDAU_JACOBIAN_COUNT, 103 LANDAU_JACOBIAN, 104 LANDAU_MASS, 105 LANDAU_F_DF, 106 LANDAU_KERNEL, 107 KSP_FACTOR, 108 KSP_SOLVE, 109 LANDAU_NUM_TIMERS 110 } LandauOMPTimers; 111 112 typedef struct { 113 PetscBool interpolate; /* Generate intermediate mesh elements */ 114 PetscBool gpu_assembly; 115 MPI_Comm comm; /* global communicator to use for errors and diagnostics */ 116 double times[LANDAU_NUM_TIMERS]; 117 PetscBool use_matrix_mass; 118 /* FE */ 119 PetscFE fe[LANDAU_MAX_SPECIES]; 120 /* geometry */ 121 PetscReal radius[LANDAU_MAX_GRIDS]; 122 PetscReal radius_par[LANDAU_MAX_GRIDS]; 123 PetscReal radius_perp[LANDAU_MAX_GRIDS]; 124 PetscReal re_radius; /* RE: radius of refinement along v_perp=0, z>0 */ 125 PetscReal vperp0_radius1; /* RE: radius of refinement along v_perp=0 */ 126 PetscReal vperp0_radius2; /* RE: radius of refinement along v_perp=0 after origin AMR refinement */ 127 PetscBool sphere; 128 PetscReal sphere_inner_radius_90degree; 129 PetscReal sphere_inner_radius_45degree; 130 PetscInt cells0[3]; 131 /* AMR */ 132 PetscBool use_p4est; 133 PetscInt numRERefine; /* RE: refinement along v_perp=0, z > 0 */ 134 PetscInt nZRefine1; /* RE: origin refinement after v_perp=0 refinement */ 135 PetscInt nZRefine2; /* RE: origin refinement after origin AMR refinement */ 136 PetscInt numAMRRefine[LANDAU_MAX_GRIDS]; /* normal AMR - refine from origin */ 137 PetscInt postAMRRefine[LANDAU_MAX_GRIDS]; /* uniform refinement of AMR */ 138 PetscBool simplex; 139 char filename[PETSC_MAX_PATH_LEN]; 140 PetscReal thermal_speed[LANDAU_MAX_GRIDS]; 141 /* relativistic */ 142 PetscBool use_energy_tensor_trick; 143 PetscBool use_relativistic_corrections; 144 /* physics */ 145 PetscReal thermal_temps[LANDAU_MAX_SPECIES]; 146 PetscReal masses[LANDAU_MAX_SPECIES]; /* mass of each species */ 147 PetscReal charges[LANDAU_MAX_SPECIES]; /* charge of each species */ 148 PetscReal n[LANDAU_MAX_SPECIES]; /* number density of each species */ 149 PetscReal m_0; /* reference mass */ 150 PetscReal v_0; /* reference velocity */ 151 PetscReal n_0; /* reference number density */ 152 PetscReal t_0; /* reference time */ 153 PetscReal Ez; 154 PetscReal epsilon0; 155 PetscReal k; 156 PetscReal lambdas[LANDAU_MAX_GRIDS][LANDAU_MAX_GRIDS]; 157 PetscReal electronShift; 158 PetscInt num_species; 159 PetscInt num_grids; 160 PetscInt species_offset[LANDAU_MAX_GRIDS + 1]; // for each grid, but same for all batched vertices 161 PetscInt mat_offset[LANDAU_MAX_GRIDS + 1]; // for each grid, but same for all batched vertices 162 // batching 163 PetscBool jacobian_field_major_order; // this could be a type but lets not get pedantic 164 VecScatter plex_batch; 165 Vec work_vec; 166 IS batch_is; 167 PetscErrorCode (*seqaij_mult)(Mat, Vec, Vec); 168 PetscErrorCode (*seqaij_multtranspose)(Mat, Vec, Vec); 169 PetscErrorCode (*seqaij_solve)(Mat, Vec, Vec); 170 PetscErrorCode (*seqaij_getdiagonal)(Mat, Vec); 171 /* COO */ 172 Mat J; 173 Mat M; 174 Vec X; 175 /* derived type */ 176 void *data; 177 /* computing */ 178 LandauDeviceType deviceType; 179 DM pack; 180 DM plex[LANDAU_MAX_GRIDS]; 181 LandauStaticData SData_d; /* static geometric data on device */ 182 /* diagnostics */ 183 PetscInt verbose; 184 PetscLogEvent events[20]; 185 PetscLogStage stage; 186 PetscObjectState norm_state; 187 PetscInt batch_sz; 188 PetscInt batch_view_idx; 189 } LandauCtx; 190 191 #define LANDAU_SPECIES_MAJOR 192 #if !defined(LANDAU_SPECIES_MAJOR) 193 #define LAND_PACK_IDX(_b, _g) (_b * ctx->num_grids + _g) 194 #define LAND_MOFFSET(_b, _g, _nbch, _ngrid, _mat_off) (_b * _mat_off[_ngrid] + _mat_off[_g]) 195 #else 196 #define LAND_PACK_IDX(_b, _g) (_g * ctx->batch_sz + _b) 197 #define LAND_MOFFSET(_b, _g, _nbch, _ngrid, _mat_off) (_nbch * _mat_off[_g] + _b * (_mat_off[_g + 1] - _mat_off[_g])) 198 #endif 199 200 typedef struct { 201 PetscReal scale; 202 LandauIdx gid; // Landau matrix index (<10,000) 203 } pointInterpolationP4est; 204 typedef struct _lP4estVertexMaps { 205 LandauIdx (*gIdx)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND]; // #elems * LANDAU_MAX_NQND 206 LandauIdx num_elements; 207 LandauIdx num_reduced; 208 LandauIdx num_face; // (Q or Q^2 for 3D) 209 LandauDeviceType deviceType; 210 PetscInt Nf; 211 pointInterpolationP4est (*c_maps)[LANDAU_MAX_Q_FACE]; 212 struct _lP4estVertexMaps *d_self; 213 void *vp1, *vp2, *vp3; 214 PetscInt numgrids; 215 } P4estVertexMaps; 216 217 #if defined(PETSC_HAVE_KOKKOS) 218 PETSC_EXTERN PetscErrorCode LandauKokkosJacobian(DM[], const PetscInt, const PetscInt, const PetscInt, const PetscInt, const PetscInt[], PetscReal[], PetscScalar[], const PetscScalar[], const LandauStaticData *, const PetscReal, const PetscLogEvent[], const PetscInt[], const PetscInt[], Mat[], Mat); 219 PETSC_EXTERN PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps *, pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE], PetscInt[], PetscInt); 220 PETSC_EXTERN PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *, PetscInt); 221 PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataSet(DM, const PetscInt, const PetscInt, const PetscInt, const PetscInt, PetscInt[], PetscInt[], PetscInt[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], LandauStaticData *); 222 PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *); 223 #endif 224 225 #endif /* PETSCLANDAU_H */ 226