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 PetscErrorCode ZLayoutCreate(PetscMPIInt size, const PetscInt eextent[3], const PetscInt vextent[3], ZLayout *zlayout) 60 { 61 ZLayout layout; 62 63 PetscFunctionBegin; 64 layout.eextent.i = eextent[0]; 65 layout.eextent.j = eextent[1]; 66 layout.eextent.k = eextent[2]; 67 layout.vextent.i = vextent[0]; 68 layout.vextent.j = vextent[1]; 69 layout.vextent.k = vextent[2]; 70 layout.comm_size = size; 71 PetscCall(PetscMalloc1(size + 1, &layout.zstarts)); 72 73 PetscInt total_elems = eextent[0] * eextent[1] * eextent[2]; 74 ZCode z = 0; 75 layout.zstarts[0] = 0; 76 for (PetscMPIInt r = 0; r < size; r++) { 77 PetscInt elems_needed = (total_elems / size) + (total_elems % size > r), count; 78 for (count = 0; count < elems_needed; z++) { 79 Ijk loc = ZCodeSplit(z); 80 if (IjkActive(layout.eextent, loc)) count++; 81 } 82 // Pick up any extra vertices in the Z ordering before the next rank's first owned element. 83 // 84 // TODO: This leads to poorly balanced vertices when eextent is a power of 2, since all the fringe vertices end up 85 // on the last rank. A possible solution is to balance the Z-order vertices independently from the cells, which will 86 // result in a lot of element closures being remote. We could finish marking boundary conditions, then do a round of 87 // vertex ownership smoothing (which would reorder and redistribute vertices without touching element distribution). 88 // Another would be to have an analytic ownership criteria for vertices in the fringe veextent - eextent. This would 89 // complicate the job of identifying an owner and its offset. 90 for (; z <= ZEncode(layout.vextent); z++) { 91 Ijk loc = ZCodeSplit(z); 92 if (IjkActive(layout.eextent, loc)) break; 93 } 94 layout.zstarts[r + 1] = z; 95 } 96 *zlayout = layout; 97 PetscFunctionReturn(0); 98 } 99 100 PetscInt ZCodeFind(ZCode key, PetscInt n, const ZCode X[]) 101 { 102 PetscInt lo = 0, hi = n; 103 104 if (n == 0) return -1; 105 while (hi - lo > 1) { 106 PetscInt mid = lo + (hi - lo) / 2; 107 if (key < X[mid]) hi = mid; 108 else lo = mid; 109 } 110 return key == X[lo] ? lo : -(lo + (key > X[lo]) + 1); 111 } 112 113 PetscErrorCode DMPlexCreateBoxMesh_Tensor_SFC_Internal(DM dm, PetscInt dim, const PetscInt faces[], const PetscReal lower[], const PetscReal upper[], const DMBoundaryType periodicity[], PetscBool interpolate) 114 { 115 PetscInt eextent[3] = {1, 1, 1}, vextent[3] = {1, 1, 1}; 116 const Ijk closure_1[] = { 117 {0, 0, 0}, 118 {1, 0, 0}, 119 }; 120 const Ijk closure_2[] = { 121 {0, 0, 0}, 122 {1, 0, 0}, 123 {1, 1, 0}, 124 {0, 1, 0}, 125 }; 126 const Ijk closure_3[] = { 127 {0, 0, 0}, 128 {0, 1, 0}, 129 {1, 1, 0}, 130 {1, 0, 0}, 131 {0, 0, 1}, 132 {1, 0, 1}, 133 {1, 1, 1}, 134 {0, 1, 1}, 135 }; 136 const Ijk *const closure_dim[] = {NULL, closure_1, closure_2, closure_3}; 137 // This must be kept consistent with DMPlexCreateCubeMesh_Internal 138 const PetscInt face_marker_1[] = {1, 2}; 139 const PetscInt face_marker_2[] = {4, 2, 1, 3}; 140 const PetscInt face_marker_3[] = {6, 5, 3, 4, 1, 2}; 141 const PetscInt *const face_marker_dim[] = {NULL, face_marker_1, face_marker_2, face_marker_3}; 142 143 PetscFunctionBegin; 144 PetscValidPointer(dm, 1); 145 PetscValidLogicalCollectiveInt(dm, dim, 2); 146 PetscCall(DMSetDimension(dm, dim)); 147 PetscMPIInt rank, size; 148 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)dm), &size)); 149 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank)); 150 for (PetscInt i = 0; i < dim; i++) { 151 eextent[i] = faces[i]; 152 vextent[i] = faces[i] + 1; 153 } 154 ZLayout layout; 155 PetscCall(ZLayoutCreate(size, eextent, vextent, &layout)); 156 PetscZSet vset; // set of all vertices in the closure of the owned elements 157 PetscCall(PetscZSetCreate(&vset)); 158 PetscInt local_elems = 0; 159 for (ZCode z = layout.zstarts[rank]; z < layout.zstarts[rank + 1]; z++) { 160 Ijk loc = ZCodeSplit(z); 161 if (IjkActive(layout.vextent, loc)) PetscZSetAdd(vset, z); 162 if (IjkActive(layout.eextent, loc)) { 163 local_elems++; 164 // Add all neighboring vertices to set 165 for (PetscInt n = 0; n < PetscPowInt(2, dim); n++) { 166 Ijk inc = closure_dim[dim][n]; 167 Ijk nloc = {loc.i + inc.i, loc.j + inc.j, loc.k + inc.k}; 168 ZCode v = ZEncode(nloc); 169 PetscZSetAdd(vset, v); 170 } 171 } 172 } 173 PetscInt local_verts, off = 0; 174 ZCode *vert_z; 175 PetscCall(PetscZSetGetSize(vset, &local_verts)); 176 PetscCall(PetscMalloc1(local_verts, &vert_z)); 177 PetscCall(PetscZSetGetElems(vset, &off, vert_z)); 178 PetscCall(PetscZSetDestroy(&vset)); 179 // ZCode is unsigned for bitwise convenience, but highest bit should never be set, so can interpret as signed 180 PetscCall(PetscSortInt64(local_verts, (PetscInt64 *)vert_z)); 181 182 PetscCall(DMPlexSetChart(dm, 0, local_elems + local_verts)); 183 for (PetscInt e = 0; e < local_elems; e++) PetscCall(DMPlexSetConeSize(dm, e, PetscPowInt(2, dim))); 184 PetscCall(DMSetUp(dm)); 185 { 186 PetscInt e = 0; 187 for (ZCode z = layout.zstarts[rank]; z < layout.zstarts[rank + 1]; z++) { 188 Ijk loc = ZCodeSplit(z); 189 if (!IjkActive(layout.eextent, loc)) continue; 190 PetscInt cone[8], orient[8] = {0}; 191 for (PetscInt n = 0; n < PetscPowInt(2, dim); n++) { 192 Ijk inc = closure_dim[dim][n]; 193 Ijk nloc = {loc.i + inc.i, loc.j + inc.j, loc.k + inc.k}; 194 ZCode v = ZEncode(nloc); 195 PetscInt ci = ZCodeFind(v, local_verts, vert_z); 196 PetscAssert(ci >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Did not find neighbor vertex in set"); 197 cone[n] = local_elems + ci; 198 } 199 PetscCall(DMPlexSetCone(dm, e, cone)); 200 PetscCall(DMPlexSetConeOrientation(dm, e, orient)); 201 e++; 202 } 203 } 204 205 if (0) { 206 DMLabel depth; 207 PetscCall(DMCreateLabel(dm, "depth")); 208 PetscCall(DMPlexGetDepthLabel(dm, &depth)); 209 PetscCall(DMLabelSetStratumBounds(depth, 0, local_elems, local_elems + local_verts)); 210 PetscCall(DMLabelSetStratumBounds(depth, 1, 0, local_elems)); 211 } else { 212 PetscCall(DMPlexSymmetrize(dm)); 213 PetscCall(DMPlexStratify(dm)); 214 } 215 { // Create point SF 216 PetscSF sf; 217 PetscSFCreate(PetscObjectComm((PetscObject)dm), &sf); 218 PetscInt owned_verts = ZCodeFind(layout.zstarts[rank + 1], local_verts, vert_z); 219 if (owned_verts < 0) owned_verts = -(owned_verts + 1); // We don't care whether the key was found 220 PetscInt num_ghosts = local_verts - owned_verts; // Due to sorting, owned vertices always come first 221 PetscInt *local_ghosts; 222 PetscSFNode *ghosts; 223 PetscCall(PetscMalloc1(num_ghosts, &local_ghosts)); 224 PetscCall(PetscMalloc1(num_ghosts, &ghosts)); 225 for (PetscInt i = 0; i < num_ghosts;) { 226 ZCode z = vert_z[owned_verts + i]; 227 PetscInt remote_rank = ZCodeFind(z, size + 1, layout.zstarts), remote_count = 0; 228 if (remote_rank < 0) remote_rank = -(remote_rank + 1) - 1; 229 // We have a new remote rank; find all the ghost indices (which are contiguous in vert_z) 230 231 // Count the elements on remote_rank 232 PetscInt remote_elem = 0; 233 for (ZCode rz = layout.zstarts[remote_rank]; rz < layout.zstarts[remote_rank + 1]; rz++) { 234 Ijk loc = ZCodeSplit(rz); 235 if (IjkActive(layout.eextent, loc)) remote_elem++; 236 } 237 238 // Traverse vertices and make ghost links 239 for (ZCode rz = layout.zstarts[remote_rank]; rz < layout.zstarts[remote_rank + 1]; rz++) { 240 Ijk loc = ZCodeSplit(rz); 241 if (rz == z) { 242 local_ghosts[i] = local_elems + owned_verts + i; 243 ghosts[i].rank = remote_rank; 244 ghosts[i].index = remote_elem + remote_count; 245 i++; 246 if (i == num_ghosts) break; 247 z = vert_z[owned_verts + i]; 248 } 249 if (IjkActive(layout.vextent, loc)) remote_count++; 250 } 251 } 252 PetscCall(PetscSFSetGraph(sf, local_elems + local_verts, num_ghosts, local_ghosts, PETSC_OWN_POINTER, ghosts, PETSC_OWN_POINTER)); 253 PetscCall(PetscObjectSetName((PetscObject)sf, "SFC Point SF")); 254 PetscCall(DMSetPointSF(dm, sf)); 255 PetscCall(PetscSFDestroy(&sf)); 256 } 257 { 258 Vec coordinates; 259 PetscScalar *coords; 260 PetscSection coord_section; 261 PetscInt coord_size; 262 PetscCall(DMGetCoordinateSection(dm, &coord_section)); 263 PetscCall(PetscSectionSetNumFields(coord_section, 1)); 264 PetscCall(PetscSectionSetFieldComponents(coord_section, 0, dim)); 265 PetscCall(PetscSectionSetChart(coord_section, local_elems, local_elems + local_verts)); 266 for (PetscInt v = 0; v < local_verts; v++) { 267 PetscInt point = local_elems + v; 268 PetscCall(PetscSectionSetDof(coord_section, point, dim)); 269 PetscCall(PetscSectionSetFieldDof(coord_section, point, 0, dim)); 270 } 271 PetscCall(PetscSectionSetUp(coord_section)); 272 PetscCall(PetscSectionGetStorageSize(coord_section, &coord_size)); 273 PetscCall(VecCreate(PETSC_COMM_SELF, &coordinates)); 274 PetscCall(PetscObjectSetName((PetscObject)coordinates, "coordinates")); 275 PetscCall(VecSetSizes(coordinates, coord_size, PETSC_DETERMINE)); 276 PetscCall(VecSetBlockSize(coordinates, dim)); 277 PetscCall(VecSetType(coordinates, VECSTANDARD)); 278 PetscCall(VecGetArray(coordinates, &coords)); 279 for (PetscInt v = 0; v < local_verts; v++) { 280 Ijk loc = ZCodeSplit(vert_z[v]); 281 coords[v * dim + 0] = lower[0] + loc.i * (upper[0] - lower[0]) / layout.eextent.i; 282 if (dim > 1) coords[v * dim + 1] = lower[1] + loc.j * (upper[1] - lower[1]) / layout.eextent.j; 283 if (dim > 2) coords[v * dim + 2] = lower[2] + loc.k * (upper[2] - lower[2]) / layout.eextent.k; 284 } 285 PetscCall(VecRestoreArray(coordinates, &coords)); 286 PetscCall(DMSetCoordinatesLocal(dm, coordinates)); 287 PetscCall(VecDestroy(&coordinates)); 288 PetscCall(PetscFree(layout.zstarts)); 289 } 290 if (interpolate) { 291 PetscCall(DMPlexInterpolateInPlace_Internal(dm)); 292 293 DMLabel label; 294 PetscCall(DMCreateLabel(dm, "Face Sets")); 295 PetscCall(DMGetLabel(dm, "Face Sets", &label)); 296 PetscInt fStart, fEnd, vStart, vEnd; 297 PetscCall(DMPlexGetHeightStratum(dm, 1, &fStart, &fEnd)); 298 PetscCall(DMPlexGetDepthStratum(dm, 0, &vStart, &vEnd)); 299 for (PetscInt f = fStart; f < fEnd; f++) { 300 PetscInt npoints; 301 PetscInt *points = NULL; 302 PetscCall(DMPlexGetTransitiveClosure(dm, f, PETSC_TRUE, &npoints, &points)); 303 PetscInt bc_count[6] = {0}; 304 for (PetscInt i = 0; i < npoints; i++) { 305 PetscInt p = points[2 * i]; 306 if (p < vStart || vEnd <= p) continue; 307 Ijk loc = ZCodeSplit(vert_z[p - vStart]); 308 // Convention here matches DMPlexCreateCubeMesh_Internal 309 bc_count[0] += loc.i == 0; 310 bc_count[1] += loc.i == layout.vextent.i - 1; 311 bc_count[2] += loc.j == 0; 312 bc_count[3] += loc.j == layout.vextent.j - 1; 313 bc_count[4] += loc.k == 0; 314 bc_count[5] += loc.k == layout.vextent.k - 1; 315 } 316 PetscCall(DMPlexRestoreTransitiveClosure(dm, f, PETSC_TRUE, &npoints, &points)); 317 for (PetscInt bc = 0, bc_match = 0; bc < 2 * dim; bc++) { 318 if (bc_count[bc] == PetscPowInt(2, dim - 1)) { 319 PetscAssert(bc_match == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Face matches multiple face sets"); 320 PetscCall(DMLabelSetValue(label, f, face_marker_dim[dim][bc])); 321 bc_match++; 322 } 323 } 324 } 325 // Ensure that the Coordinate DM has our new boundary labels 326 DM cdm; 327 PetscCall(DMGetCoordinateDM(dm, &cdm)); 328 PetscCall(DMCopyLabels(dm, cdm, PETSC_COPY_VALUES, PETSC_FALSE, DM_COPY_LABELS_FAIL)); 329 } 330 PetscCall(PetscFree(vert_z)); 331 PetscFunctionReturn(0); 332 } 333