Date: 2022-08-02T03:15:10+00:00

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 */