From 6239d2576489a3178d753471f988908f5499577c Mon Sep 17 00:00:00 2001 From: Weiqun Zhang Date: Fri, 4 Oct 2024 20:48:27 -0500 Subject: [PATCH] New function for domain decomposition (#4183) Add a new function that decomposition a Box into BoxArray. The number of Boxes in the returned BoxArray matches the given nboxes argument, unless the domain is too small. We aim to decompose the domain into subdomains that are as cubic as possible, even if this results in Boxes with odd numbers of cells. Thus, this function is generally not suited for applications with multiple AMR levels or for multigrid solvers. However, it could be useful for single-level applications that prefer exactly one Box per process. --- Src/Base/AMReX_BoxArray.H | 18 ++++ Src/Base/AMReX_BoxArray.cpp | 170 ++++++++++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+) diff --git a/Src/Base/AMReX_BoxArray.H b/Src/Base/AMReX_BoxArray.H index b3b339c33b..e85946872c 100644 --- a/Src/Base/AMReX_BoxArray.H +++ b/Src/Base/AMReX_BoxArray.H @@ -53,6 +53,24 @@ namespace amrex //! Note that two BoxArrays that match are not necessarily equal. [[nodiscard]] bool match (const BoxArray& x, const BoxArray& y); + /** + * \brief Decompose domain box into BoxArray + * + * The returned BoxArray has nboxes Boxes, unless the the domain is too + * small. We aim to decompose the domain into subdomains that are as + * cubic as possible, even if this results in Boxes with odd numbers of + * cells. Thus, this function is generally not suited for applications + * with multiple AMR levels or for multigrid solvers. + * + * \param domain Domain Box + * \param nboxes the target number of Boxes + * \param decomp controls whether domain decomposition should be done in + * that direction. + */ + [[nodiscard]] BoxArray decompose (Box const& domain, int nboxes, + Array const& decomp + = {AMREX_D_DECL(true,true,true)}); + struct BARef { BARef (); diff --git a/Src/Base/AMReX_BoxArray.cpp b/Src/Base/AMReX_BoxArray.cpp index 9bca594352..576d4cb870 100644 --- a/Src/Base/AMReX_BoxArray.cpp +++ b/Src/Base/AMReX_BoxArray.cpp @@ -12,6 +12,9 @@ #include +#include +#include +#include #include namespace amrex { @@ -1887,6 +1890,173 @@ bool match (const BoxArray& x, const BoxArray& y) } } +BoxArray decompose (Box const& domain, int nboxes, + Array const& decomp) +{ + auto ndecomp = std::count(decomp.begin(), decomp.end(), true); + + if (nboxes <= 1 || ndecomp == 0) { + return BoxArray(domain); + } + + Box const& ccdomain = amrex::enclosedCells(domain); + IntVect const& ncells = ccdomain.length(); + IntVect nprocs(1); + + if (ndecomp == 1) { + for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) { + if (decomp[idim]) { + nprocs[idim] = nboxes; + } + } + } else { + // Factorization of nboxes + Vector factors; + { + int x = 2; + int n = nboxes; + while (x*x <= n) { + std::div_t dv = std::div(n, x); + if (dv.rem == 0) { + factors.push_back(x); + n = dv.quot; + } else { + ++x; + } + } + if (n != 1) { + factors.push_back(n); + } + AMREX_ALWAYS_ASSERT(nboxes == std::accumulate(factors.begin(), factors.end(), + 1, std::multiplies<>())); + } + + struct ProcDim + { + int nproc; + int idim; + Vector procs; + ProcDim (int np, int dim) : nproc(np), idim(dim) {} + }; + + Vector procdim; + procdim.reserve(AMREX_SPACEDIM); + + Array nblocks; + + for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) { + if (decomp[idim]) { + nblocks[idim] = ncells[idim]; + procdim.emplace_back(1,idim); + } else { + nblocks[idim] = 0; // This dimension will not be decomposed. + } + } + + auto comp = [&] (ProcDim const& a, ProcDim const& b) { + if (nblocks[a.idim]*b.nproc < + nblocks[b.idim]*a.nproc) { + return true; + } else if (nblocks[a.idim]*b.nproc > + nblocks[b.idim]*a.nproc) { + return false; + } else { + return a.procs.size() > b.procs.size(); + } + }; + + int nprocs_tot = 1; + while (!factors.empty()) { + std::sort(procdim.begin(), procdim.end(), comp); + auto f = factors.back(); + factors.pop_back(); + procdim.back().nproc *= f; + procdim.back().procs.push_back(f); + nprocs_tot *= f; + if (nprocs_tot == nboxes) { + break; + } + } + + // swap to see if the decomposition can be improved. + while (true) + { + std::sort(procdim.begin(), procdim.end(), comp); + auto fit = std::find_if(procdim.begin(),procdim.end(), + [] (ProcDim const& x) { return x.nproc > 1; }); + if (fit == procdim.end()) { break; } // This should not actually happen. + auto& light = *fit; + auto& heavy = procdim.back(); + Long w0 = nblocks[light.idim] * heavy.nproc; + Long w1 = nblocks[heavy.idim] * light.nproc; + if (w0 >= w1) { break; } + bool swapped = false; + for (auto& f0 : light.procs) { + for (auto& f1 : heavy.procs) { + if ((f0 > f1) && (w0*f0 < w1*f1)) { + light.nproc /= f0; + light.nproc *= f1; + heavy.nproc /= f1; + heavy.nproc *= f0; + std::swap(f0,f1); + swapped = true; + break; + } + } + if (swapped) { break;} + } + if (!swapped) { break; } + } + + for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) { + if (!decomp[idim]) { + procdim.emplace_back(1,idim); + } + } + for (auto const& pd : procdim) { + nprocs[pd.idim] = pd.nproc; + } + } + + AMREX_ALWAYS_ASSERT(AMREX_D_TERM(nprocs[0],*nprocs[1],*nprocs[2]) == nboxes); + + IntVect const domlo = ccdomain.smallEnd(); + IntVect const sz = ncells / nprocs; + IntVect const extra = ncells - sz*nprocs; + auto ixtyp = domain.ixType(); + BoxList bl(ixtyp); +#if (AMREX_SPACEDIM == 3) + for (int k = 0; k < nprocs[2]; ++k) { + // The first extra[2] blocks get one extra cell with a total of + // sz[2]+1. The rest get sz[2] cells. The decomposition in y + // and x directions are similar. + int klo = (k < extra[2]) ? k*(sz[2]+1) : (k*sz[2]+extra[2]); + int khi = (k < extra[2]) ? klo+(sz[2]+1)-1 : klo+sz[2]-1; + klo += domlo[2]; + khi += domlo[2]; +#endif +#if (AMREX_SPACEDIM >= 2) + for (int j = 0; j < nprocs[1]; ++j) { + int jlo = (j < extra[1]) ? j*(sz[1]+1) : (j*sz[1]+extra[1]); + int jhi = (j < extra[1]) ? jlo+(sz[1]+1)-1 : jlo+sz[1]-1; + jlo += domlo[1]; + jhi += domlo[1]; +#endif + for (int i = 0; i < nprocs[0]; ++i) { + int ilo = (i < extra[0]) ? i*(sz[0]+1) : (i*sz[0]+extra[0]); + int ihi = (i < extra[0]) ? ilo+(sz[0]+1)-1 : ilo+sz[0]-1; + ilo += domlo[0]; + ihi += domlo[0]; + Box b{IntVect(AMREX_D_DECL(ilo,jlo,klo)), + IntVect(AMREX_D_DECL(ihi,jhi,khi))}; + if (b.ok()) { + bl.push_back(b.convert(ixtyp)); + } + AMREX_D_TERM(},},}) + + return BoxArray(std::move(bl)); +} + std::ostream& operator<< (std::ostream& os, const BoxArray::RefID& id) {