diff --git a/Src/Base/AMReX_BoxArray.H b/Src/Base/AMReX_BoxArray.H index b3b339c33b..0850a17ac2 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); + /** + * \breif 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..3b385bed0d 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,170 @@ bool match (const BoxArray& x, const BoxArray& y) } } +BoxArray decompose (Box const& domain, int nboxes, + Array const& decomp) +{ + int 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& light = procdim.front(); + 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) {