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 *const closure_dim[] = {NULL, closure_1, closure_2, closure_3}; 127 // This must be kept consistent with DMPlexCreateCubeMesh_Internal 128 const PetscInt face_marker_1[] = {1, 2}; 129 const PetscInt face_marker_2[] = {4, 2, 1, 3}; 130 const PetscInt face_marker_3[] = {6, 5, 3, 4, 1, 2}; 131 const PetscInt *const face_marker_dim[] = {NULL, face_marker_1, face_marker_2, face_marker_3}; 132 133 PetscFunctionBegin; 134 PetscValidPointer(dm, 1); 135 PetscValidLogicalCollectiveInt(dm, dim, 2); 136 PetscCall(DMSetDimension(dm, dim)); 137 PetscMPIInt rank, size; 138 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)dm), &size)); 139 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank)); 140 for (PetscInt i = 0; i < dim; i++) { 141 eextent[i] = faces[i]; 142 vextent[i] = faces[i] + 1; 143 } 144 ZLayout layout = ZLayoutCreate(size, eextent, vextent); 145 PetscZSet vset; // set of all vertices in the closure of the owned elements 146 PetscCall(PetscZSetCreate(&vset)); 147 PetscInt local_elems = 0; 148 for (ZCode z = layout.zstarts[rank]; z < layout.zstarts[rank + 1]; z++) { 149 Ijk loc = ZCodeSplit(z); 150 if (IjkActive(layout.vextent, loc)) PetscZSetAdd(vset, z); 151 if (IjkActive(layout.eextent, loc)) { 152 local_elems++; 153 // Add all neighboring vertices to set 154 for (PetscInt n = 0; n < PetscPowInt(2, dim); n++) { 155 Ijk inc = closure_dim[dim][n]; 156 Ijk nloc = {loc.i + inc.i, loc.j + inc.j, loc.k + inc.k}; 157 ZCode v = ZEncode(nloc); 158 PetscZSetAdd(vset, v); 159 } 160 } 161 } 162 PetscInt local_verts, off = 0; 163 ZCode *vert_z; 164 PetscCall(PetscZSetGetSize(vset, &local_verts)); 165 PetscCall(PetscMalloc1(local_verts, &vert_z)); 166 PetscCall(PetscZSetGetElems(vset, &off, vert_z)); 167 PetscCall(PetscZSetDestroy(&vset)); 168 // ZCode is unsigned for bitwise convenience, but highest bit should never be set, so can interpret as signed 169 PetscCall(PetscSortInt64(local_verts, (PetscInt64 *)vert_z)); 170 171 PetscCall(DMPlexSetChart(dm, 0, local_elems + local_verts)); 172 for (PetscInt e = 0; e < local_elems; e++) PetscCall(DMPlexSetConeSize(dm, e, PetscPowInt(2, dim))); 173 PetscCall(DMSetUp(dm)); 174 { 175 PetscInt e = 0; 176 for (ZCode z = layout.zstarts[rank]; z < layout.zstarts[rank + 1]; z++) { 177 Ijk loc = ZCodeSplit(z); 178 if (!IjkActive(layout.eextent, loc)) continue; 179 PetscInt cone[8], orient[8] = {0}; 180 for (PetscInt n = 0; n < PetscPowInt(2, dim); n++) { 181 Ijk inc = closure_dim[dim][n]; 182 Ijk nloc = {loc.i + inc.i, loc.j + inc.j, loc.k + inc.k}; 183 ZCode v = ZEncode(nloc); 184 PetscInt ci = ZCodeFind(v, local_verts, vert_z); 185 PetscAssert(ci >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Did not find neighbor vertex in set"); 186 cone[n] = local_elems + ci; 187 } 188 PetscCall(DMPlexSetCone(dm, e, cone)); 189 PetscCall(DMPlexSetConeOrientation(dm, e, orient)); 190 e++; 191 } 192 } 193 194 if (0) { 195 DMLabel depth; 196 PetscCall(DMCreateLabel(dm, "depth")); 197 PetscCall(DMPlexGetDepthLabel(dm, &depth)); 198 PetscCall(DMLabelSetStratumBounds(depth, 0, local_elems, local_elems + local_verts)); 199 PetscCall(DMLabelSetStratumBounds(depth, 1, 0, local_elems)); 200 } else { 201 PetscCall(DMPlexSymmetrize(dm)); 202 PetscCall(DMPlexStratify(dm)); 203 } 204 { // Create point SF 205 PetscSF sf; 206 PetscSFCreate(PetscObjectComm((PetscObject)dm), &sf); 207 PetscInt owned_verts = ZCodeFind(layout.zstarts[rank + 1], local_verts, vert_z); 208 if (owned_verts < 0) owned_verts = -(owned_verts + 1); // We don't care whether the key was found 209 PetscInt num_ghosts = local_verts - owned_verts; // Due to sorting, owned vertices always come first 210 PetscInt *local_ghosts; 211 PetscSFNode *ghosts; 212 PetscCall(PetscMalloc1(num_ghosts, &local_ghosts)); 213 PetscCall(PetscMalloc1(num_ghosts, &ghosts)); 214 for (PetscInt i = 0; i < num_ghosts;) { 215 ZCode z = vert_z[owned_verts + i]; 216 PetscInt remote_rank = ZCodeFind(z, size + 1, layout.zstarts), remote_count = 0; 217 if (remote_rank < 0) remote_rank = -(remote_rank + 1) - 1; 218 // We have a new remote rank; find all the ghost indices (which are contiguous in vert_z) 219 220 // Count the elements on remote_rank 221 PetscInt remote_elem = 0; 222 for (ZCode rz = layout.zstarts[remote_rank]; rz < layout.zstarts[remote_rank + 1]; rz++) { 223 Ijk loc = ZCodeSplit(rz); 224 if (IjkActive(layout.eextent, loc)) remote_elem++; 225 } 226 227 // Traverse vertices and make ghost links 228 for (ZCode rz = layout.zstarts[remote_rank]; rz < layout.zstarts[remote_rank + 1]; rz++) { 229 Ijk loc = ZCodeSplit(rz); 230 if (rz == z) { 231 local_ghosts[i] = local_elems + owned_verts + i; 232 ghosts[i].rank = remote_rank; 233 ghosts[i].index = remote_elem + remote_count; 234 i++; 235 if (i == num_ghosts) break; 236 z = vert_z[owned_verts + i]; 237 } 238 if (IjkActive(layout.vextent, loc)) remote_count++; 239 } 240 } 241 PetscCall(PetscSFSetGraph(sf, local_elems + local_verts, num_ghosts, local_ghosts, PETSC_OWN_POINTER, ghosts, PETSC_OWN_POINTER)); 242 PetscCall(PetscObjectSetName((PetscObject)sf, "SFC Point SF")); 243 PetscCall(DMSetPointSF(dm, sf)); 244 PetscCall(PetscSFDestroy(&sf)); 245 } 246 { 247 Vec coordinates; 248 PetscScalar *coords; 249 PetscSection coord_section; 250 PetscInt coord_size; 251 PetscCall(DMGetCoordinateSection(dm, &coord_section)); 252 PetscCall(PetscSectionSetNumFields(coord_section, 1)); 253 PetscCall(PetscSectionSetFieldComponents(coord_section, 0, dim)); 254 PetscCall(PetscSectionSetChart(coord_section, local_elems, local_elems + local_verts)); 255 for (PetscInt v = 0; v < local_verts; v++) { 256 PetscInt point = local_elems + v; 257 PetscCall(PetscSectionSetDof(coord_section, point, dim)); 258 PetscCall(PetscSectionSetFieldDof(coord_section, point, 0, dim)); 259 } 260 PetscCall(PetscSectionSetUp(coord_section)); 261 PetscCall(PetscSectionGetStorageSize(coord_section, &coord_size)); 262 PetscCall(VecCreate(PETSC_COMM_SELF, &coordinates)); 263 PetscCall(PetscObjectSetName((PetscObject)coordinates, "coordinates")); 264 PetscCall(VecSetSizes(coordinates, coord_size, PETSC_DETERMINE)); 265 PetscCall(VecSetBlockSize(coordinates, dim)); 266 PetscCall(VecSetType(coordinates, VECSTANDARD)); 267 PetscCall(VecGetArray(coordinates, &coords)); 268 for (PetscInt v = 0; v < local_verts; v++) { 269 Ijk loc = ZCodeSplit(vert_z[v]); 270 coords[v * dim + 0] = lower[0] + loc.i * (upper[0] - lower[0]) / layout.eextent.i; 271 if (dim > 1) coords[v * dim + 1] = lower[1] + loc.j * (upper[1] - lower[1]) / layout.eextent.j; 272 if (dim > 2) coords[v * dim + 2] = lower[2] + loc.k * (upper[2] - lower[2]) / layout.eextent.k; 273 } 274 PetscCall(VecRestoreArray(coordinates, &coords)); 275 PetscCall(DMSetCoordinatesLocal(dm, coordinates)); 276 PetscCall(VecDestroy(&coordinates)); 277 PetscCall(PetscFree(layout.zstarts)); 278 } 279 if (interpolate) { 280 PetscCall(DMPlexInterpolateInPlace_Internal(dm)); 281 282 DMLabel label; 283 PetscCall(DMCreateLabel(dm, "Face Sets")); 284 PetscCall(DMGetLabel(dm, "Face Sets", &label)); 285 PetscInt fStart, fEnd, vStart, vEnd; 286 PetscCall(DMPlexGetHeightStratum(dm, 1, &fStart, &fEnd)); 287 PetscCall(DMPlexGetDepthStratum(dm, 0, &vStart, &vEnd)); 288 for (PetscInt f = fStart; f < fEnd; f++) { 289 PetscInt npoints; 290 PetscInt *points = NULL; 291 PetscCall(DMPlexGetTransitiveClosure(dm, f, PETSC_TRUE, &npoints, &points)); 292 PetscInt bc_count[6] = {0}; 293 for (PetscInt i = 0; i < npoints; i++) { 294 PetscInt p = points[2 * i]; 295 if (p < vStart || vEnd <= p) continue; 296 Ijk loc = ZCodeSplit(vert_z[p - vStart]); 297 // Convention here matches DMPlexCreateCubeMesh_Internal 298 bc_count[0] += loc.i == 0; 299 bc_count[1] += loc.i == layout.vextent.i - 1; 300 bc_count[2] += loc.j == 0; 301 bc_count[3] += loc.j == layout.vextent.j - 1; 302 bc_count[4] += loc.k == 0; 303 bc_count[5] += loc.k == layout.vextent.k - 1; 304 } 305 PetscCall(DMPlexRestoreTransitiveClosure(dm, f, PETSC_TRUE, &npoints, &points)); 306 for (PetscInt bc = 0, bc_match = 0; bc < 2 * dim; bc++) { 307 if (bc_count[bc] == PetscPowInt(2, dim - 1)) { 308 PetscAssert(bc_match == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Face matches multiple face sets"); 309 PetscCall(DMLabelSetValue(label, f, face_marker_dim[dim][bc])); 310 bc_match++; 311 } 312 } 313 } 314 // Ensure that the Coordinate DM has our new boundary labels 315 DM cdm; 316 PetscCall(DMGetCoordinateDM(dm, &cdm)); 317 PetscCall(DMCopyLabels(dm, cdm, PETSC_COPY_VALUES, PETSC_FALSE, DM_COPY_LABELS_FAIL)); 318 } 319 PetscCall(PetscFree(vert_z)); 320 PetscFunctionReturn(0); 321 } 322