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