Actual source code: petsclandau.h

  1: #if !defined(PETSCLANDAU_H)
  2: #define PETSCLANDAU_H

  4: #include <petscdmplex.h>
  5: #include <petscts.h>

  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 DMPlexLandauAddMaxwellians(DM, Vec, PetscReal, PetscReal[], PetscReal[], PetscInt, PetscInt, void *);
 11: PETSC_EXTERN PetscErrorCode DMPlexLandauCreateMassMatrix(DM dm, Mat *Amat);
 12: PETSC_EXTERN PetscErrorCode DMPlexLandauIFunction(TS, PetscReal,Vec,Vec,Vec,void *);
 13: PETSC_EXTERN PetscErrorCode DMPlexLandauIJacobian(TS, PetscReal,Vec,Vec,PetscReal,Mat,Mat,void *);

 15: typedef int LandauIdx;

 17: /* the Fokker-Planck-Landau context */
 18: #if !defined(LANDAU_DIM)
 19: #define LANDAU_DIM 2
 20: #endif

 22: #if !defined(LANDAU_MAX_SPECIES)
 23: #if LANDAU_DIM==2
 24: #define LANDAU_MAX_SPECIES 10
 25: #define LANDAU_MAX_GRIDS 3
 26: #else
 27: #define LANDAU_MAX_SPECIES 3
 28: #define LANDAU_MAX_GRIDS 2
 29: #endif
 30: #else
 31: #define LANDAU_MAX_GRIDS 3
 32: #endif

 34: #if !defined(LANDAU_MAX_Q)
 35: #if defined(LANDAU_MAX_NQ)
 36: #error"LANDAU_MAX_NQ but not LANDAU_MAX_Q. Use -DLANDAU_MAX_Q=4 for Q3 elements"
 37: #endif
 38: #if LANDAU_DIM==2
 39: #define LANDAU_MAX_Q 5
 40: #else
 41: #define LANDAU_MAX_Q 3
 42: #endif
 43: #else
 44: #undef LANDAU_MAX_NQ
 45: #endif

 47: #if LANDAU_DIM==2
 48: #define LANDAU_MAX_Q_FACE LANDAU_MAX_Q
 49: #define LANDAU_MAX_NQ (LANDAU_MAX_Q*LANDAU_MAX_Q)
 50: #define LANDAU_MAX_BATCH_SZ 320
 51: #else
 52: #define LANDAU_MAX_Q_FACE (LANDAU_MAX_Q*LANDAU_MAX_Q)
 53: #define LANDAU_MAX_NQ (LANDAU_MAX_Q*LANDAU_MAX_Q*LANDAU_MAX_Q)
 54: #define LANDAU_MAX_BATCH_SZ 32
 55: #endif

 57: typedef enum {LANDAU_CUDA, LANDAU_KOKKOS, LANDAU_CPU} LandauDeviceType;

 59: // static data - will be "device" data
 60: typedef struct {
 61:   void  *invJ;  // nip*dim*dim
 62:   void  *D;     // nq*nb*dim
 63:   void  *B;     // nq*nb
 64:   void  *alpha; // ns
 65:   void  *beta;  // ns
 66:   void  *invMass; // ns
 67:   void  *w; // nip
 68:   void  *x; // nip
 69:   void  *y; // nip
 70:   void  *z; // nip
 71:   void  *Eq_m; // ns - dynamic
 72:   void  *f; //  nip*Nf - dynamic (IP)
 73:   void  *dfdx; // nip*Nf - dynamic (IP)
 74:   void  *dfdy; // nip*Nf - dynamic (IP)
 75:   void  *dfdz; // nip*Nf - dynamic (IP)
 76:   int   dim_,ns_,nip_,nq_,nb_;
 77:   void  *NCells; // remove and ise elem_offset - TODO
 78:   void  *species_offset; // for each grid, but same for all batched vertices
 79:   void  *mat_offset; // for each grid, but same for all batched vertices
 80:   void  *elem_offset; // for each grid, but same for all batched vertices
 81:   void  *ip_offset; // for each grid, but same for all batched vertices
 82:   void  *ipf_offset; // for each grid, but same for all batched vertices
 83:   void  *ipfdf_data; // for each grid, but same for all batched vertices
 84:   void  *maps; // for each grid, but same for all batched vertices
 85:   // COO
 86:   void             *coo_elem_offsets;
 87:   void             *coo_elem_point_offsets;
 88:   void             *coo_elem_fullNb;
 89:   LandauIdx        coo_n_cellsTot;
 90:   LandauIdx        coo_size;
 91:   LandauIdx        coo_max_fullnb;
 92: } LandauStaticData;

 94: typedef enum {LANDAU_EX2_TSSOLVE, LANDAU_MATRIX_TOTAL, LANDAU_OPERATOR, LANDAU_JACOBIAN_COUNT, LANDAU_JACOBIAN, LANDAU_MASS, LANDAU_F_DF, LANDAU_KERNEL, KSP_FACTOR, KSP_SOLVE, LANDAU_NUM_TIMERS} LandauOMPTimers;

 96: typedef struct {
 97:   PetscBool      interpolate;                  /* Generate intermediate mesh elements */
 98:   PetscBool      gpu_assembly;
 99:   MPI_Comm       comm; /* global communicator to use for errors and diagnostics */
100:   double         times[LANDAU_NUM_TIMERS];
101:   PetscBool      use_matrix_mass;
102:   /* FE */
103:   PetscFE        fe[LANDAU_MAX_SPECIES];
104:   /* geometry  */
105:   PetscReal      radius[LANDAU_MAX_GRIDS];
106:   PetscReal      re_radius;           /* RE: radius of refinement along v_perp=0, z>0 */
107:   PetscReal      vperp0_radius1;      /* RE: radius of refinement along v_perp=0 */
108:   PetscReal      vperp0_radius2;      /* RE: radius of refinement along v_perp=0 after origin AMR refinement */
109:   PetscBool      sphere; // not used
110:   PetscBool      inflate; // not used
111:   PetscReal      i_radius[LANDAU_MAX_GRIDS]; // not used
112:   PetscReal      e_radius; // not used
113:   PetscInt       num_sections; // not used
114:   /* AMR */
115:   PetscBool      use_p4est;
116:   PetscInt       numRERefine;        /* RE: refinement along v_perp=0, z > 0 */
117:   PetscInt       nZRefine1;          /* RE: origin refinement after v_perp=0 refinement */
118:   PetscInt       nZRefine2;          /* RE: origin refinement after origin AMR refinement */
119:   PetscInt       numAMRRefine[LANDAU_MAX_GRIDS];  /* normal AMR - refine from origin */
120:   PetscInt       postAMRRefine[LANDAU_MAX_GRIDS]; /* uniform refinement of AMR */
121:   /* relativistic */
122:   PetscBool      use_energy_tensor_trick;
123:   PetscBool      use_relativistic_corrections;
124:   /* physics */
125:   PetscReal      thermal_temps[LANDAU_MAX_SPECIES];
126:   PetscReal      masses[LANDAU_MAX_SPECIES];  /* mass of each species  */
127:   PetscReal      charges[LANDAU_MAX_SPECIES]; /* charge of each species  */
128:   PetscReal      n[LANDAU_MAX_SPECIES];       /* number density of each species  */
129:   PetscReal      m_0;      /* reference mass */
130:   PetscReal      v_0;      /* reference velocity */
131:   PetscReal      n_0;      /* reference number density */
132:   PetscReal      t_0;      /* reference time */
133:   PetscReal      Ez;
134:   PetscReal      epsilon0;
135:   PetscReal      k;
136:   PetscReal      lnLam;
137:   PetscReal      electronShift;
138:   PetscInt       num_species;
139:   PetscInt       num_grids;
140:   PetscInt       species_offset[LANDAU_MAX_GRIDS+1]; // for each grid, but same for all batched vertices
141:   PetscInt       mat_offset[LANDAU_MAX_GRIDS+1]; // for each grid, but same for all batched vertices
142:   // batching
143:   PetscBool      jacobian_field_major_order; // this could be a type but lets not get pedantic
144:   VecScatter     plex_batch;
145:   Vec            work_vec;
146:   IS             batch_is;
147:   PetscErrorCode (*seqaij_mult)(Mat,Vec,Vec);
148:   PetscErrorCode (*seqaij_multtranspose)(Mat,Vec,Vec);
149:   PetscErrorCode (*seqaij_solve)(Mat,Vec,Vec);
150:   PetscErrorCode (*seqaij_getdiagonal)(Mat,Vec);
151:   /* COO */
152:   PetscBool        coo_assembly;
153:   /* cache */
154:   Mat              J;
155:   Mat              M;
156:   Vec              X;
157:   /* derived type */
158:   void             *data;
159:   /* computing */
160:   LandauDeviceType deviceType;
161:   DM               pack;
162:   DM               plex[LANDAU_MAX_GRIDS];
163:   LandauStaticData SData_d; /* static geometric data on device */
164:   /* diagnostics */
165:   PetscInt         verbose;
166:   PetscLogEvent    events[20];
167:   PetscLogStage    stage;
168:   PetscBool        aux_bool;  /* helper */
169:   PetscInt         batch_sz;
170:   PetscInt         batch_view_idx;
171: } LandauCtx;

173: #define LANDAU_SPECIES_MAJOR
174: #if !defined(LANDAU_SPECIES_MAJOR)
175: #define LAND_PACK_IDX(_b,_g) (_b*ctx->num_grids + _g)
176: #define LAND_MOFFSET(_b,_g,_nbch,_ngrid,_mat_off) (_b*_mat_off[_ngrid] + _mat_off[_g])
177: #else
178: #define LAND_PACK_IDX(_b,_g) (_g*ctx->batch_sz + _b)
179: #define LAND_MOFFSET(_b,_g,_nbch,_ngrid,_mat_off) (_nbch*_mat_off[_g] + _b*(_mat_off[_g+1] - _mat_off[_g]))
180: #endif

182: typedef struct {
183:   PetscReal scale;
184:   LandauIdx gid;   // Landau matrix index (<10,000)
185: } pointInterpolationP4est;
186: typedef struct _lP4estVertexMaps {
187:   LandauIdx                (*gIdx)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ]; // #elems *  LANDAU_MAX_NQ (spoof for max , Nb) on device,
188:   LandauIdx                num_elements;
189:   LandauIdx                num_reduced;
190:   LandauIdx                num_face;  // (Q or Q^2 for 3D)
191:   LandauDeviceType         deviceType;
192:   PetscInt                 Nf;
193:   pointInterpolationP4est (*c_maps)[LANDAU_MAX_Q_FACE];
194:   struct _lP4estVertexMaps*d_self;
195:   void                    *vp1,*vp2,*vp3;
196:   PetscInt                numgrids;
197: } P4estVertexMaps;

199: PETSC_EXTERN PetscErrorCode LandauCreateColoring(Mat, DM, PetscContainer *);
200: #if defined(PETSC_HAVE_CUDA)
201: PETSC_EXTERN PetscErrorCode LandauCUDAJacobian(DM[], const PetscInt, const PetscInt, const PetscInt, const PetscInt[], PetscReal[], PetscScalar[], const PetscScalar[],
202:                                                const LandauStaticData *, const PetscReal, const PetscLogEvent[],const PetscInt[], const PetscInt[], Mat[], Mat);
203: PETSC_EXTERN PetscErrorCode LandauCUDACreateMatMaps(P4estVertexMaps *, pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE], PetscInt[], PetscInt, PetscInt);
204: PETSC_EXTERN PetscErrorCode LandauCUDADestroyMatMaps(P4estVertexMaps *, PetscInt);
205: PETSC_EXTERN PetscErrorCode LandauCUDAStaticDataSet(DM, const PetscInt, const PetscInt, const PetscInt, PetscInt[], PetscInt[], PetscInt[], PetscReal [], PetscReal [], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[],
206:                                                     PetscReal[], LandauStaticData *);
207: PETSC_EXTERN PetscErrorCode LandauCUDAStaticDataClear(LandauStaticData *);
208: #endif
209: #if defined(PETSC_HAVE_KOKKOS)
210: PETSC_EXTERN PetscErrorCode LandauKokkosJacobian(DM[], const PetscInt, const PetscInt, const PetscInt, const PetscInt[], PetscReal[], PetscScalar[],  const PetscScalar[],
211:                                                  const LandauStaticData *, const PetscReal, const PetscLogEvent[], const PetscInt[], const PetscInt[], Mat[], Mat);
212: PETSC_EXTERN PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps *, pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE], PetscInt[], PetscInt, PetscInt);
213: PETSC_EXTERN PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *, PetscInt);
214: PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataSet(DM, const PetscInt, const PetscInt, const PetscInt, PetscInt[], PetscInt[], PetscInt[], PetscReal [], PetscReal [], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[],
215:                                                       PetscReal[], LandauStaticData *);
216: PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData*);
217: #endif

219: #endif /* PETSCLANDAU_H */