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