1 #include <petsc/private/dmpleximpl.h> /*I "petscdmplex.h" I*/ 2 #include <petscsf.h> 3 #include <petsc/private/hashset.h> 4 5 typedef uint64_t ZCode; 6 7 PETSC_HASH_SET(ZSet, ZCode, PetscHash_UInt64, PetscHashEqual) 8 9 typedef struct { 10 PetscInt i, j, k; 11 } Ijk; 12 13 typedef struct { 14 Ijk eextent; 15 Ijk vextent; 16 PetscMPIInt comm_size; 17 ZCode *zstarts; 18 } ZLayout; 19 20 static unsigned ZCodeSplit1(ZCode z) 21 { 22 z = ((z & 01001001001001001) | ((z >> 2) & 02002002002002002) | ((z >> 4) & 04004004004004004)); 23 z = (z | (z >> 6) | (z >> 12)) & 0000000777000000777; 24 z = (z | (z >> 18)) & 0777777; 25 return (unsigned)z; 26 } 27 28 static ZCode ZEncode1(unsigned t) 29 { 30 ZCode z = t; 31 z = (z | (z << 18)) & 0777000000777; 32 z = (z | (z << 6) | (z << 12)) & 07007007007007007; 33 z = (z | (z << 2) | (z << 4)) & 0111111111111111111; 34 return z; 35 } 36 37 static Ijk ZCodeSplit(ZCode z) 38 { 39 Ijk c; 40 c.i = ZCodeSplit1(z >> 2); 41 c.j = ZCodeSplit1(z >> 1); 42 c.k = ZCodeSplit1(z >> 0); 43 return c; 44 } 45 46 static ZCode ZEncode(Ijk c) 47 { 48 ZCode z = (ZEncode1(c.i) << 2) | (ZEncode1(c.j) << 1) | ZEncode1(c.k); 49 return z; 50 } 51 52 static PetscBool IjkActive(Ijk extent, Ijk l) 53 { 54 if (l.i < extent.i && l.j < extent.j && l.k < extent.k) return PETSC_TRUE; 55 return PETSC_FALSE; 56 } 57 58 // Since element/vertex box extents are typically not equal powers of 2, Z codes that lie within the domain are not contiguous. 59 static ZLayout ZLayoutCreate(PetscMPIInt size, const PetscInt eextent[3], const PetscInt vextent[3]) 60 { 61 ZLayout layout; 62 layout.eextent.i = eextent[0]; 63 layout.eextent.j = eextent[1]; 64 layout.eextent.k = eextent[2]; 65 layout.vextent.i = vextent[0]; 66 layout.vextent.j = vextent[1]; 67 layout.vextent.k = vextent[2]; 68 layout.comm_size = size; 69 PetscMalloc1(size + 1, &layout.zstarts); 70 71 PetscInt total_elems = eextent[0] * eextent[1] * eextent[2]; 72 ZCode z = 0; 73 layout.zstarts[0] = 0; 74 for (PetscMPIInt r = 0; r < size; r++) { 75 PetscInt elems_needed = (total_elems / size) + (total_elems % size > r), count; 76 for (count = 0; count < elems_needed; z++) { 77 Ijk loc = ZCodeSplit(z); 78 if (IjkActive(layout.eextent, loc)) count++; 79 } 80 // Pick up any extra vertices in the Z ordering before the next rank's first owned element. 81 for (; z <= ZEncode(layout.vextent); z++) { 82 Ijk loc = ZCodeSplit(z); 83 if (IjkActive(layout.eextent, loc)) break; 84 } 85 layout.zstarts[r + 1] = z; 86 } 87 return layout; 88 } 89 90 PetscInt ZCodeFind(ZCode key, PetscInt n, const ZCode X[]) 91 { 92 PetscInt lo = 0, hi = n; 93 94 if (n == 0) return -1; 95 while (hi - lo > 1) { 96 PetscInt mid = lo + (hi - lo) / 2; 97 if (key < X[mid]) hi = mid; 98 else lo = mid; 99 } 100 return key == X[lo] ? lo : -(lo + (key > X[lo]) + 1); 101 } 102 103 PetscErrorCode DMPlexCreateBoxMesh_Tensor_SFC_Internal(DM dm, PetscInt dim, const PetscInt faces[], const PetscReal lower[], const PetscReal upper[], const DMBoundaryType periodicity[], PetscBool interpolate) 104 { 105 PetscInt eextent[3] = {1, 1, 1}, vextent[3] = {1, 1, 1}; 106 const Ijk closure_1[] = { 107 {0, 0, 0}, 108 {1, 0, 0}, 109 }; 110 const Ijk closure_2[] = { 111 {0, 0, 0}, 112 {1, 0, 0}, 113 {1, 1, 0}, 114 {0, 1, 0}, 115 }; 116 const Ijk closure_3[] = { 117 {0, 0, 0}, 118 {0, 1, 0}, 119 {1, 1, 0}, 120 {1, 0, 0}, 121 {0, 0, 1}, 122 {1, 0, 1}, 123 {1, 1, 1}, 124 {0, 1, 1}, 125 }; 126 const Ijk *closure_dim[] = {NULL, closure_1, closure_2, closure_3}; 127 128 PetscFunctionBegin; 129 PetscValidPointer(dm, 1); 130 PetscValidLogicalCollectiveInt(dm, dim, 2); 131 PetscCall(DMSetDimension(dm, dim)); 132 PetscMPIInt rank, size; 133 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)dm), &size)); 134 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank)); 135 for (PetscInt i = 0; i < dim; i++) { 136 eextent[i] = faces[i]; 137 vextent[i] = faces[i] + 1; 138 } 139 ZLayout layout = ZLayoutCreate(size, eextent, vextent); 140 PetscZSet vset; // set of all vertices in the closure of the owned elements 141 PetscCall(PetscZSetCreate(&vset)); 142 PetscInt local_elems = 0; 143 for (ZCode z = layout.zstarts[rank]; z < layout.zstarts[rank + 1]; z++) { 144 Ijk loc = ZCodeSplit(z); 145 if (IjkActive(layout.vextent, loc)) PetscZSetAdd(vset, z); 146 if (IjkActive(layout.eextent, loc)) { 147 local_elems++; 148 // Add all neighboring vertices to set 149 for (PetscInt n = 0; n < PetscPowInt(2, dim); n++) { 150 Ijk inc = closure_dim[dim][n]; 151 Ijk nloc = {loc.i + inc.i, loc.j + inc.j, loc.k + inc.k}; 152 ZCode v = ZEncode(nloc); 153 PetscZSetAdd(vset, v); 154 } 155 } 156 } 157 PetscInt local_verts, off = 0; 158 ZCode *vert_z; 159 PetscCall(PetscZSetGetSize(vset, &local_verts)); 160 PetscCall(PetscMalloc1(local_verts, &vert_z)); 161 PetscCall(PetscZSetGetElems(vset, &off, vert_z)); 162 PetscCall(PetscZSetDestroy(&vset)); 163 // ZCode is unsigned for bitwise convenience, but highest bit should never be set, so can interpret as signed 164 PetscCall(PetscSortInt64(local_verts, (PetscInt64 *)vert_z)); 165 166 PetscCall(DMPlexSetChart(dm, 0, local_elems + local_verts)); 167 for (PetscInt e = 0; e < local_elems; e++) PetscCall(DMPlexSetConeSize(dm, e, PetscPowInt(2, dim))); 168 PetscCall(DMSetUp(dm)); 169 { 170 PetscInt e = 0; 171 for (ZCode z = layout.zstarts[rank]; z < layout.zstarts[rank + 1]; z++) { 172 Ijk loc = ZCodeSplit(z); 173 if (!IjkActive(layout.eextent, loc)) continue; 174 PetscInt cone[8], orient[8] = {0}; 175 for (PetscInt n = 0; n < PetscPowInt(2, dim); n++) { 176 Ijk inc = closure_dim[dim][n]; 177 Ijk nloc = {loc.i + inc.i, loc.j + inc.j, loc.k + inc.k}; 178 ZCode v = ZEncode(nloc); 179 PetscInt ci = ZCodeFind(v, local_verts, vert_z); 180 PetscAssert(ci >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Did not find neighbor vertex in set"); 181 cone[n] = local_elems + ci; 182 } 183 PetscCall(DMPlexSetCone(dm, e, cone)); 184 PetscCall(DMPlexSetConeOrientation(dm, e, orient)); 185 e++; 186 } 187 } 188 189 if (0) { 190 DMLabel depth; 191 PetscCall(DMCreateLabel(dm, "depth")); 192 PetscCall(DMPlexGetDepthLabel(dm, &depth)); 193 PetscCall(DMLabelSetStratumBounds(depth, 0, local_elems, local_elems + local_verts)); 194 PetscCall(DMLabelSetStratumBounds(depth, 1, 0, local_elems)); 195 } else { 196 PetscCall(DMPlexSymmetrize(dm)); 197 PetscCall(DMPlexStratify(dm)); 198 } 199 { // Create point SF 200 PetscSF sf; 201 PetscSFCreate(PetscObjectComm((PetscObject)dm), &sf); 202 PetscInt owned_verts = ZCodeFind(layout.zstarts[rank + 1], local_verts, vert_z); 203 if (owned_verts < 0) owned_verts = -(owned_verts + 1); // We don't care whether the key was found 204 PetscInt num_ghosts = local_verts - owned_verts; // Due to sorting, owned vertices always come first 205 PetscInt *local_ghosts; 206 PetscSFNode *ghosts; 207 PetscCall(PetscMalloc1(num_ghosts, &local_ghosts)); 208 PetscCall(PetscMalloc1(num_ghosts, &ghosts)); 209 for (PetscInt i = 0; i < num_ghosts;) { 210 ZCode z = vert_z[owned_verts + i]; 211 PetscInt remote_rank = ZCodeFind(z, size + 1, layout.zstarts), remote_count = 0; 212 if (remote_rank < 0) remote_rank = -(remote_rank + 1) - 1; 213 // We have a new remote rank; find all the ghost indices (which are contiguous in vert_z) 214 215 // Count the elements on remote_rank 216 PetscInt remote_elem = 0; 217 for (ZCode rz = layout.zstarts[remote_rank]; rz < layout.zstarts[remote_rank + 1]; rz++) { 218 Ijk loc = ZCodeSplit(rz); 219 if (IjkActive(layout.eextent, loc)) remote_elem++; 220 } 221 222 // Traverse vertices and make ghost links 223 for (ZCode rz = layout.zstarts[remote_rank]; rz < layout.zstarts[remote_rank + 1]; rz++) { 224 Ijk loc = ZCodeSplit(rz); 225 if (rz == z) { 226 local_ghosts[i] = local_elems + owned_verts + i; 227 ghosts[i].rank = remote_rank; 228 ghosts[i].index = remote_elem + remote_count; 229 i++; 230 if (i == num_ghosts) break; 231 z = vert_z[owned_verts + i]; 232 } 233 if (IjkActive(layout.vextent, loc)) remote_count++; 234 } 235 } 236 PetscCall(PetscSFSetGraph(sf, local_elems + local_verts, num_ghosts, local_ghosts, PETSC_OWN_POINTER, ghosts, PETSC_OWN_POINTER)); 237 PetscCall(PetscObjectSetName((PetscObject)sf, "SFC Point SF")); 238 PetscCall(DMSetPointSF(dm, sf)); 239 PetscCall(PetscSFDestroy(&sf)); 240 } 241 242 { 243 Vec coordinates; 244 PetscScalar *coords; 245 PetscSection coord_section; 246 PetscInt coord_size; 247 DM cdm; 248 PetscCall(DMGetCoordinateSection(dm, &coord_section)); 249 PetscCall(PetscSectionSetNumFields(coord_section, 1)); 250 PetscCall(PetscSectionSetFieldComponents(coord_section, 0, dim)); 251 PetscCall(PetscSectionSetChart(coord_section, local_elems, local_elems + local_verts)); 252 for (PetscInt v = 0; v < local_verts; v++) { 253 PetscInt point = local_elems + v; 254 PetscCall(PetscSectionSetDof(coord_section, point, dim)); 255 PetscCall(PetscSectionSetFieldDof(coord_section, point, 0, dim)); 256 } 257 PetscCall(PetscSectionSetUp(coord_section)); 258 PetscCall(PetscSectionGetStorageSize(coord_section, &coord_size)); 259 PetscCall(VecCreate(PETSC_COMM_SELF, &coordinates)); 260 PetscCall(PetscObjectSetName((PetscObject)coordinates, "coordinates")); 261 PetscCall(VecSetSizes(coordinates, coord_size, PETSC_DETERMINE)); 262 PetscCall(VecSetBlockSize(coordinates, dim)); 263 PetscCall(VecSetType(coordinates, VECSTANDARD)); 264 PetscCall(VecGetArray(coordinates, &coords)); 265 for (PetscInt v = 0; v < local_verts; v++) { 266 Ijk loc = ZCodeSplit(vert_z[v]); 267 coords[v * dim + 0] = lower[0] + loc.i * (upper[0] - lower[0]) / layout.eextent.i; 268 if (dim > 1) coords[v * dim + 1] = lower[1] + loc.j * (upper[1] - lower[1]) / layout.eextent.j; 269 if (dim > 2) coords[v * dim + 2] = lower[2] + loc.k * (upper[2] - lower[2]) / layout.eextent.k; 270 } 271 PetscCall(VecRestoreArray(coordinates, &coords)); 272 PetscCall(DMSetCoordinatesLocal(dm, coordinates)); 273 PetscCall(VecDestroy(&coordinates)); 274 PetscCall(PetscFree(vert_z)); 275 PetscCall(PetscFree(layout.zstarts)); 276 277 PetscFE fe; 278 PetscCall(DMGetCoordinateDM(dm, &cdm)); 279 PetscCall(PetscFECreateLagrange(PetscObjectComm((PetscObject)dm), dim, dim, PETSC_FALSE, 1, PETSC_DETERMINE, &fe)); 280 PetscCall(DMSetField(cdm, 0, NULL, (PetscObject)fe)); 281 PetscCall(PetscFEDestroy(&fe)); 282 PetscCall(DMCreateDS(cdm)); 283 } 284 if (interpolate) { 285 DM dmint; 286 PetscCall(DMPlexInterpolate(dm, &dmint)); 287 PetscCall(DMPlexReplace_Internal(dm, &dmint)); 288 } 289 PetscFunctionReturn(0); 290 } 291