diff --git a/src/utils/index_split.cpp b/src/utils/index_split.cpp index bbd6ec11bf653..814907f096ffd 100644 --- a/src/utils/index_split.cpp +++ b/src/utils/index_split.cpp @@ -26,103 +26,112 @@ #include "mesh/mesh.hpp" namespace parthenon { - struct DummyFunctor { DummyFunctor() = default; KOKKOS_INLINE_FUNCTION void operator()(team_mbr_t team_member) const {} }; -IndexSplit::IndexSplit(MeshData *md, const IndexRange &kb, const IndexRange &jb, - const IndexRange &ib, const int nkp, const int njp) - : nghost_(Globals::nghost), nkp_(nkp), njp_(njp), kbs_(kb.s), jbs_(jb.s), ibs_(ib.s), - ibe_(ib.e) { - Init(md, kb.e, jb.e); - ndim_ = md->GetNDim(); -} +IndexSplit::IndexSplit(MeshData *md, IndexDomain domain, const int nk_tiles, + const int nj_tiles, TopologicalElement te, + TopologicalElement te_mem) + : IndexSplit(md, md->GetBoundsK(domain, te), md->GetBoundsJ(domain, te), + md->GetBoundsI(domain, te), nk_tiles, nj_tiles, te_mem) {} -IndexSplit::IndexSplit(MeshData *md, IndexDomain domain, const int nkp, - const int njp) - : nghost_(Globals::nghost), nkp_(nkp), njp_(njp) { - auto ib = md->GetBoundsI(domain); - auto jb = md->GetBoundsJ(domain); - auto kb = md->GetBoundsK(domain); - kbs_ = kb.s; - jbs_ = jb.s; - ibs_ = ib.s; - ibe_ = ib.e; - Init(md, kb.e, jb.e); - ndim_ = md->GetNDim(); +IndexSplit IndexSplit::RawMemIJ(IndexDomain domain, int halo, MeshData *md, + TE logical_te, TE memory_te) { + auto pmesh = md->GetMeshPointer(); + const int ndim = pmesh->ndim; + auto ib = md->GetBoundsI(domain, logical_te); + auto jb = md->GetBoundsJ(domain, logical_te); + auto kb = md->GetBoundsK(domain, logical_te); + ib.s -= halo; + ib.e += halo; + jb.s -= (ndim > 1) * halo; + jb.e += (ndim > 1) * halo; + kb.s -= (ndim > 2) * halo; + kb.e += (ndim > 2) * halo; + int nk_tiles = kb.e - kb.s + 1; // Outer loop iterates over all k + int nj_tiles = + 1; // Tile contains all j indices, so inner loops run over all i and j for a fixed k + return IndexSplit(md, kb, jb, ib, kb.e - kb.s + 1, 1); } -void IndexSplit::Init(MeshData *md, const int kbe, const int jbe) { - const int total_k = kbe - kbs_ + 1; - const int total_j = jbe - jbs_ + 1; - const int total_i = ibe_ - ibs_ + 1; +IndexSplit::IndexSplit(MeshData *md, const IndexRange &kb, const IndexRange &jb, + const IndexRange &ib, const int nk_tiles, const int nj_tiles, + TopologicalElement te_mem) + : nk_tiles_(nk_tiles), nj_tiles_(nj_tiles) { + // nk_tiles_ and nj_tiles_ define how the kj space is tiled into (nk_tiles_ x nj_tiles_) + // tiles. The k- and j- bounds of each of the tiles are returned by `GetBoundsK` and + // `GetBoundsJ`. The loop structure is: + // - Outermost loop over tiles + // - Middle loop over k range of the tile (since that can't be pulled into the inner + // contiguous memory loop) + // - Inner contiguous memory loop over i-range and j-range of tile, including ghosts + // where necessary + + // Save the size of the logical domain (i.e. the requested index range) + logical_ = Indexer3D(kb, jb, ib); + + // save the size of the memory domain of the block we are iterating over + using TE = TopologicalElement; + PARTHENON_REQUIRE( + te_mem == TE::CC || te_mem == TE::NN, + "All memory layouts either are cell-centered or nodal, even for faces and edges."); + auto mib = md->GetBoundsI(IndexDomain::entire, te_mem); + auto mjb = md->GetBoundsJ(IndexDomain::entire, te_mem); + auto mkb = md->GetBoundsK(IndexDomain::entire, te_mem); + + memory_ = Indexer3D(mkb, mjb, mib); // Compute max parallelism (at outer loop level) from Kokkos // equivalent to NSMS in Kokkos // TODO(JMM): I'm not sure if this is really the best way to do // this. Based on discussion on Kokkos slack. + int concurrency{1}; // = NSMs = 132 for NVIDIA H100 #ifdef PARTHENON_ENABLE_GPU const auto space = DevExecSpace(); - team_policy policy(space, (md->NumBlocks()) * total_k, Kokkos::AUTO); + team_policy policy(space, (md->NumBlocks()) * logical_.Extent(), Kokkos::AUTO); // JMM: In principle, should pass a realistic functor here. Using a // dummy because we don't know what's available. // TODO(JMM): Should we expose the functor? - policy.set_scratch_size(1, Kokkos::PerTeam(sizeof(Real) * total_i * total_j)); + policy.set_scratch_size(1, Kokkos::PerTeam(sizeof(Real) * logical_.Extent() * + logical_.Extent())); const int nteams = policy.team_size_recommended(DummyFunctor(), Kokkos::ParallelForTag()); - concurrency_ = space.concurrency() / nteams; -#else - concurrency_ = 1; + concurrency = space.concurrency() / nteams; #endif // PARTHENON_ENABLE_GPU - if (nkp_ == all_outer) - nkp_ = total_k; - else if (nkp_ == no_outer) - nkp_ = 1; - if (njp_ == all_outer) - njp_ = total_j; - else if (njp_ == no_outer) - njp_ = 1; + if (nk_tiles_ == all_outer) + nk_tiles_ = logical_.Extent(); + else if (nk_tiles_ == no_outer) + nk_tiles_ = 1; + if (nj_tiles_ == all_outer) + nj_tiles_ = logical_.Extent(); + else if (nj_tiles_ == no_outer) + nj_tiles_ = 1; - if (nkp_ == 0) { + if (nk_tiles_ == 0) { #ifdef PARTHENON_ENABLE_GPU - nkp_ = total_k; + nk_tiles_ = logical_.Extent(); #else - nkp_ = 1; + nk_tiles_ = 1; #endif // PARTHENON_ENABLE_GPU - } else if (nkp_ > total_k) { - nkp_ = total_k; + } else if (nk_tiles_ > logical_.Extent()) { + nk_tiles_ = logical_.Extent(); } - if (njp_ == 0) { + if (nj_tiles_ == 0) { #ifdef PARTHENON_ENABLE_GPU // From Forrest Glines: - // nkp_ * njp_ >= number of SMs / number of streams - // => njp_ >= SMS / streams / NKP - njp_ = std::min(concurrency_ / (NSTREAMS_ * nkp_), total_j); + // nk_tiles_ * nj_tiles_ >= number of SMs / number of streams + // => nj_tiles_ >= SMS / streams / nk_tiles + nj_tiles_ = std::min(concurrency / (NSTREAMS_ * nk_tiles_), logical_.Extent()); #else - njp_ = 1; + nj_tiles_ = 1; #endif // PARTHENON_ENABLE_GPU - } else if (njp_ > total_j) { - njp_ = total_j; + } else if (nj_tiles_ > logical_.Extent()) { + nj_tiles_ = logical_.Extent(); } - - // add a tiny bit to avoid round-off issues when we ultimately convert to int - // JMM: Do NOT cast these to integers here. The casting happens later. - // These being doubles is necessary for proper interleaving of work. - target_k_ = (1.0 * total_k) / nkp_ + 1.e-6; - target_j_ = (1.0 * total_j) / njp_ + 1.e-6; - - // save the "entire" ranges - // don't bother save ".s" since it's always zero - auto ib = md->GetBoundsI(IndexDomain::entire); - auto jb = md->GetBoundsJ(IndexDomain::entire); - auto kb = md->GetBoundsK(IndexDomain::entire); - kbe_entire_ = kb.e; - jbe_entire_ = jb.e; - ibe_entire_ = ib.e; } } // namespace parthenon diff --git a/src/utils/index_split.hpp b/src/utils/index_split.hpp index 84b71210766f3..c8e47ca761758 100644 --- a/src/utils/index_split.hpp +++ b/src/utils/index_split.hpp @@ -14,6 +14,9 @@ #ifndef UTILS_INDEX_SPLIT_HPP_ #define UTILS_INDEX_SPLIT_HPP_ +#include +#include + #include "basic_types.hpp" #include "defs.hpp" #include "globals.hpp" @@ -26,101 +29,188 @@ template class MeshData; class IndexSplit { + using TE = TopologicalElement; + public: static constexpr int all_outer = -100; static constexpr int no_outer = -200; IndexSplit(MeshData *md, const IndexRange &kb, const IndexRange &jb, - const IndexRange &ib, const int nkp, const int njp); - IndexSplit(MeshData *md, IndexDomain domain, const int nkp, const int njp); + const IndexRange &ib, const int nk_tiles, const int nj_tiles, + TE te_mem = TE::CC); + IndexSplit(MeshData *md, IndexDomain domain, const int nk_tiles, + const int nj_tiles, TE te = TE::CC, TE te_mem = TE::CC); + + // Provides the same functionality as the raw memory indexer did when initialized with + // the IJ factory + static IndexSplit RawMemIJ(IndexDomain domain, int halo, MeshData *md, + TE logical_te, TE memory_te = TE::CC); + static IndexSplit RawMemIJ(IndexDomain domain, MeshData *md, TE logical_te, + TE memory_te = TE::CC) { + return RawMemIJ(domain, 0, md, logical_te, memory_te); + } + + // Get the total number of kj-tiles + int outer_size() const { return nk_tiles_ * nj_tiles_; } + + // Temporary backward compatibility with RawMemoryIndexer + KOKKOS_INLINE_FUNCTION + auto GetStartIndices(int outer_idx) const { + PARTHENON_REQUIRE(nj_tiles_ == 1 && nk_tiles_ == logical_.Extent(), + "Only works for this case."); + auto kb = GetBoundsK(outer_idx); + auto jb = GetBoundsJ(outer_idx); + return std::tuple{kb.s, jb.s, logical_.StartIdx()}; + } + + KOKKOS_INLINE_FUNCTION + int GetNinnerRaw(int outer_idx) const { + PARTHENON_REQUIRE(nj_tiles_ == 1 && nk_tiles_ == logical_.Extent(), + "Only works for this case."); + auto [ks, js, is] = GetStartIndices(outer_idx); + int ke = ks; // Enforce fixed k by hand + int je = logical_.EndIdx(); // Enforce end of j-range by hand + int ie = logical_.EndIdx(); // Enforce end of i-range by hand + return inner_size({ks, ke}, {js, je}, {is, ie}); + } - int outer_size() const { return nkp_ * njp_; } + KOKKOS_INLINE_FUNCTION + int GetNouter() const { return outer_size(); } + + int GetMaxNinnerRaw() const { + int max_ninner{0}; + for (int p = 0; p < outer_size(); ++p) + max_ninner = std::max(max_ninner, GetNinnerRaw(p)); + return max_ninner; + } + + KOKKOS_INLINE_FUNCTION + int GetStartingRawFlatIdx(int outer_idx) const { + auto [ks, js, is] = GetStartIndices(outer_idx); + return memory_.GetFlatIdx(ks, js, is); + } + + KOKKOS_INLINE_FUNCTION + auto GetCurrentIndices(int starting_raw_flat_idx, int inner_idx) const { + return memory_(starting_raw_flat_idx + inner_idx); + } + + // Get the k-bounds of kj-tile indexed by p KOKKOS_INLINE_FUNCTION IndexRange GetBoundsK(const int p) const { - const auto kf = p / njp_; - return {kbs_ + static_cast(kf * target_k_), - kbs_ + static_cast((kf + 1) * target_k_) - 1}; + const auto k_tile = p / nj_tiles_; + const int ks = logical_.StartIdx(); + const int nk = logical_.Extent(); + const int start = ks + (k_tile * nk) / nk_tiles_; + const int stop = ks + ((k_tile + 1) * nk) / nk_tiles_ - 1; + return {start, stop}; } + + // Get the j-bounds of kj-tile indexed by p KOKKOS_INLINE_FUNCTION IndexRange GetBoundsJ(const int p) const { - const auto jf = p % njp_; - return {jbs_ + static_cast(jf * target_j_), - jbs_ + static_cast((jf + 1) * target_j_) - 1}; + const auto j_tile = p % nj_tiles_; + const int js = logical_.StartIdx(); + const int nj = logical_.Extent(); + const int start = js + (j_tile * nj) / nj_tiles_; + const int stop = js + ((j_tile + 1) * nj) / nj_tiles_ - 1; + return {start, stop}; } + KOKKOS_INLINE_FUNCTION - IndexRange GetBoundsI() const { return {ibs_, ibe_}; } + IndexRange GetBoundsI() const { + return {logical_.StartIdx(), logical_.EndIdx()}; + } + KOKKOS_INLINE_FUNCTION IndexRange GetBoundsI(const int p) const { return GetBoundsI(); } + KOKKOS_INLINE_FUNCTION - auto GetBoundsKJI(const int p) const { + std::tuple GetBoundsKJI(const int p) const { const auto kb = GetBoundsK(p); const auto jb = GetBoundsJ(p); const auto ib = GetBoundsI(p); return std::make_tuple(kb, jb, ib); } + + template + KOKKOS_INLINE_FUNCTION void middle_for(int p, F &&f) const { + // TODO(LFR): This could be generalized to allow for switching to including part of + // k-space in the flattening. + const auto [kb, jb, ib] = GetBoundsKJI(p); + for (int k = kb.s; k <= kb.e; ++k) + f(k, jb.s, ib.s, inner_size(kb, jb, ib)); + } + + KOKKOS_INLINE_FUNCTION int inner_size(const IndexRange &kb, const IndexRange &jb, + const IndexRange &ib) const { + return memory_.GetFlatIdx(kb.e, jb.e, ib.e) - memory_.GetFlatIdx(kb.s, jb.s, ib.s) + + 1; + } + KOKKOS_INLINE_FUNCTION IndexRange GetInnerBounds(const IndexRange &jb) const { - return {ibs_, (ibe_entire_ + 1) * (jb.e - jb.s + 1) - (ibe_entire_ - ibe_) - 1}; + const int ibs = logical_.StartIdx(); + const int ibe = logical_.EndIdx(); + const int kbs = logical_.StartIdx(); + return {ibs, ibs + inner_size({kbs, kbs}, jb, {ibs, ibe}) - 1}; } + KOKKOS_INLINE_FUNCTION IndexRange GetInnerBounds(const IndexRange &jb, const IndexRange &ib) const { - return {ib.s, (ibe_entire_ + 1) * (jb.e - jb.s + 1) - (ibe_entire_ - ib.e) - 1}; + const int kbs = logical_.StartIdx(); + return {ib.s, ib.s + inner_size({kbs, kbs}, jb, ib) - 1}; + } + + KOKKOS_INLINE_FUNCTION + int GetMemoryIdx(int ks, int js, int is) const { + return memory_.GetFlatIdx(ks, js, is); } KOKKOS_FORCEINLINE_FUNCTION - int get_i(const int idx) const { return idx % (ibe_entire_ + 1); } + int get_i(const int idx) const { return idx % memory_.Extent(); } + KOKKOS_FORCEINLINE_FUNCTION - int get_deltaj(const int idx) const { return idx / (ibe_entire_ + 1); } + int get_deltaj(const int idx) const { return idx / memory_.Extent(); } KOKKOS_INLINE_FUNCTION bool is_i_ghost(const int idx) const { - const int ni = ibe_entire_ + 1; - const int i = idx % ni; - const int i_inner_size = ni - 2 * nghost_; - return (i < nghost_ || i - nghost_ >= i_inner_size); + const int i = get_i(idx); + return !logical_.IdxInRange(i); } + KOKKOS_INLINE_FUNCTION bool is_j_ghost(const int outer_idx, const int idx) const { - const int ni = ibe_entire_ + 1; - const int j = GetBoundsJ(outer_idx).s + idx / ni; - const int j_inner_size = jbe_entire_ + 1 - 2 * nghost_; - return (ndim_ > 1 && (j < nghost_ || j - nghost_ >= j_inner_size)); + const int j = GetBoundsJ(outer_idx).s + idx / memory_.Extent(); + return !logical_.IdxInRange(j); } + KOKKOS_INLINE_FUNCTION - bool is_k_ghost(const int k) const { - const int k_inner_size = kbe_entire_ + 1 - 2 * nghost_; - return (ndim_ > 2 && (k < nghost_ || k - nghost_ >= k_inner_size)); - } + bool is_k_ghost(const int k) const { return !logical_.IdxInRange(k); } + KOKKOS_INLINE_FUNCTION bool is_ghost(const int outer_idx, const int k, const int idx) const { return is_k_ghost(k) || is_j_ghost(outer_idx, idx) || is_i_ghost(idx); } KOKKOS_INLINE_FUNCTION - int get_max_ni() const { return ibe_entire_ + 1; } + int get_max_ni() const { return memory_.Extent(); } // TODO(@jdolence) these overestimate max size...should probably fix KOKKOS_INLINE_FUNCTION - int get_max_nj() const { return (jbe_entire_ + 1) / njp_ + 1; } + int get_max_nj() const { return memory_.Extent() / nj_tiles_ + 1; } KOKKOS_INLINE_FUNCTION - int get_max_nk() const { return (kbe_entire_ + 1) / nkp_ + 1; } + int get_max_nk() const { return memory_.Extent() / nk_tiles_ + 1; } KOKKOS_INLINE_FUNCTION int get_max_nij() const { return get_max_ni() * get_max_nj(); } - // inner_size could be used to find the bounds for a loop that is collapsed over - // 1, 2, or 3 dimensions by providing the right starting and stopping indices - template - KOKKOS_INLINE_FUNCTION int inner_size(const V &v, const IndexRange &kb, - const IndexRange &jb, - const IndexRange &ib) const { - return &v(0, kb.e, jb.e, ib.e) - &v(0, kb.s, jb.s, ib.s); - } private: // TODO(JMM): Replace this with a macro or something when available static constexpr int NSTREAMS_ = 1; // Change if we add streams back - int concurrency_; // = NSMs = 132 for NVIDIA H100 - int nghost_, nkp_, njp_, kbs_, jbs_, ibs_, ibe_; - int kbe_entire_, jbe_entire_, ibe_entire_, ndim_; - float target_k_, target_j_; - void Init(MeshData *md, const int kbe, const int jbe); + int nk_tiles_, nj_tiles_; + Indexer3D logical_, memory_; + + static constexpr std::size_t IDIM{2}; + static constexpr std::size_t JDIM{1}; + static constexpr std::size_t KDIM{0}; }; } // namespace parthenon diff --git a/src/utils/indexer.hpp b/src/utils/indexer.hpp index 723c090ac25f0..a6c9315212310 100644 --- a/src/utils/indexer.hpp +++ b/src/utils/indexer.hpp @@ -123,6 +123,16 @@ struct Indexer { return end; } + template + KOKKOS_FORCEINLINE_FUNCTION auto IdxInRange(int i) const { + return i >= StartIdx() && i <= EndIdx(); + } + + template + KOKKOS_FORCEINLINE_FUNCTION auto Extent() const { + return N[I] / GetN(); + } + KOKKOS_FORCEINLINE_FUNCTION auto End() const { return End_impl(std::make_index_sequence()); } diff --git a/tst/unit/gold_files/index_split/2D_ng2_allk_njp3.gold b/tst/unit/gold_files/index_split/2D_ng2_allk_njp3.gold new file mode 100644 index 0000000000000..2db38a0364976 --- /dev/null +++ b/tst/unit/gold_files/index_split/2D_ng2_allk_njp3.gold @@ -0,0 +1,7 @@ +# IndexSplit Gold File: 2D_ng2_allk_njp3 +# Config: ndim=2 nx=6,6,0 nghost=2 nkp=-100 njp=3 +# Domain bounds: k=[0,0] j=[2,7] i=[2,7] +outer_size=3 +p=0 k=[0,0] j=[2,3] i=[2,7] inner=[2,17] inner_size=16 +p=1 k=[0,0] j=[4,5] i=[2,7] inner=[2,17] inner_size=16 +p=2 k=[0,0] j=[6,7] i=[2,7] inner=[2,17] inner_size=16 diff --git a/tst/unit/gold_files/index_split/3D_ng2_allk_allj.gold b/tst/unit/gold_files/index_split/3D_ng2_allk_allj.gold new file mode 100644 index 0000000000000..23209eb7eaad9 --- /dev/null +++ b/tst/unit/gold_files/index_split/3D_ng2_allk_allj.gold @@ -0,0 +1,20 @@ +# IndexSplit Gold File: 3D_ng2_allk_allj +# Config: ndim=3 nx=4,4,4 nghost=2 nkp=-100 njp=-100 +# Domain bounds: k=[2,5] j=[2,5] i=[2,5] +outer_size=16 +p=0 k=[2,2] j=[2,2] i=[2,5] inner=[2,5] inner_size=4 +p=1 k=[2,2] j=[3,3] i=[2,5] inner=[2,5] inner_size=4 +p=2 k=[2,2] j=[4,4] i=[2,5] inner=[2,5] inner_size=4 +p=3 k=[2,2] j=[5,5] i=[2,5] inner=[2,5] inner_size=4 +p=4 k=[3,3] j=[2,2] i=[2,5] inner=[2,5] inner_size=4 +p=5 k=[3,3] j=[3,3] i=[2,5] inner=[2,5] inner_size=4 +p=6 k=[3,3] j=[4,4] i=[2,5] inner=[2,5] inner_size=4 +p=7 k=[3,3] j=[5,5] i=[2,5] inner=[2,5] inner_size=4 +p=8 k=[4,4] j=[2,2] i=[2,5] inner=[2,5] inner_size=4 +p=9 k=[4,4] j=[3,3] i=[2,5] inner=[2,5] inner_size=4 +p=10 k=[4,4] j=[4,4] i=[2,5] inner=[2,5] inner_size=4 +p=11 k=[4,4] j=[5,5] i=[2,5] inner=[2,5] inner_size=4 +p=12 k=[5,5] j=[2,2] i=[2,5] inner=[2,5] inner_size=4 +p=13 k=[5,5] j=[3,3] i=[2,5] inner=[2,5] inner_size=4 +p=14 k=[5,5] j=[4,4] i=[2,5] inner=[2,5] inner_size=4 +p=15 k=[5,5] j=[5,5] i=[2,5] inner=[2,5] inner_size=4 diff --git a/tst/unit/gold_files/index_split/3D_ng2_allk_noj.gold b/tst/unit/gold_files/index_split/3D_ng2_allk_noj.gold new file mode 100644 index 0000000000000..91b65fde565d0 --- /dev/null +++ b/tst/unit/gold_files/index_split/3D_ng2_allk_noj.gold @@ -0,0 +1,8 @@ +# IndexSplit Gold File: 3D_ng2_allk_noj +# Config: ndim=3 nx=4,4,4 nghost=2 nkp=-100 njp=-200 +# Domain bounds: k=[2,5] j=[2,5] i=[2,5] +outer_size=4 +p=0 k=[2,2] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 +p=1 k=[3,3] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 +p=2 k=[4,4] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 +p=3 k=[5,5] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 diff --git a/tst/unit/gold_files/index_split/3D_ng2_asym_4x8x16.gold b/tst/unit/gold_files/index_split/3D_ng2_asym_4x8x16.gold new file mode 100644 index 0000000000000..cc9446f40b08b --- /dev/null +++ b/tst/unit/gold_files/index_split/3D_ng2_asym_4x8x16.gold @@ -0,0 +1,12 @@ +# IndexSplit Gold File: 3D_ng2_asym_4x8x16 +# Config: ndim=3 nx=4,8,16 nghost=2 nkp=4 njp=2 +# Domain bounds: k=[2,17] j=[2,9] i=[2,5] +outer_size=8 +p=0 k=[2,5] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 +p=1 k=[2,5] j=[6,9] i=[2,5] inner=[2,29] inner_size=28 +p=2 k=[6,9] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 +p=3 k=[6,9] j=[6,9] i=[2,5] inner=[2,29] inner_size=28 +p=4 k=[10,13] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 +p=5 k=[10,13] j=[6,9] i=[2,5] inner=[2,29] inner_size=28 +p=6 k=[14,17] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 +p=7 k=[14,17] j=[6,9] i=[2,5] inner=[2,29] inner_size=28 diff --git a/tst/unit/gold_files/index_split/3D_ng2_nkp2_njp1.gold b/tst/unit/gold_files/index_split/3D_ng2_nkp2_njp1.gold new file mode 100644 index 0000000000000..823a355c78b5d --- /dev/null +++ b/tst/unit/gold_files/index_split/3D_ng2_nkp2_njp1.gold @@ -0,0 +1,6 @@ +# IndexSplit Gold File: 3D_ng2_nkp2_njp1 +# Config: ndim=3 nx=4,4,4 nghost=2 nkp=2 njp=1 +# Domain bounds: k=[2,5] j=[2,5] i=[2,5] +outer_size=2 +p=0 k=[2,3] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 +p=1 k=[4,5] j=[2,5] i=[2,5] inner=[2,29] inner_size=28 diff --git a/tst/unit/gold_files/index_split/3D_ng2_nkp3_njp2.gold b/tst/unit/gold_files/index_split/3D_ng2_nkp3_njp2.gold new file mode 100644 index 0000000000000..2ceb5cc9e1517 --- /dev/null +++ b/tst/unit/gold_files/index_split/3D_ng2_nkp3_njp2.gold @@ -0,0 +1,10 @@ +# IndexSplit Gold File: 3D_ng2_nkp3_njp2 +# Config: ndim=3 nx=6,6,6 nghost=2 nkp=3 njp=2 +# Domain bounds: k=[2,7] j=[2,7] i=[2,7] +outer_size=6 +p=0 k=[2,3] j=[2,4] i=[2,7] inner=[2,27] inner_size=26 +p=1 k=[2,3] j=[5,7] i=[2,7] inner=[2,27] inner_size=26 +p=2 k=[4,5] j=[2,4] i=[2,7] inner=[2,27] inner_size=26 +p=3 k=[4,5] j=[5,7] i=[2,7] inner=[2,27] inner_size=26 +p=4 k=[6,7] j=[2,4] i=[2,7] inner=[2,27] inner_size=26 +p=5 k=[6,7] j=[5,7] i=[2,7] inner=[2,27] inner_size=26 diff --git a/tst/unit/test_index_split.cpp b/tst/unit/test_index_split.cpp index 6264993505b27..ef6061154a6a4 100644 --- a/tst/unit/test_index_split.cpp +++ b/tst/unit/test_index_split.cpp @@ -10,7 +10,10 @@ // license in this material to reproduce, prepare derivative works, distribute copies to // the public, perform publicly and display publicly, and to permit others to do so. //======================================================================================== +#include +#include #include +#include #include #include #include @@ -80,6 +83,42 @@ struct v5 : public parthenon::variable_names::base_t { : parthenon::variable_names::base_t(std::forward(args)...) {} static std::string name() { return "v5"; } }; +// Helper to create blocks with asymmetric dimensions +BlockList_t MakeBlockListAsymmetric(const std::shared_ptr pkg, + const int NBLOCKS, const int nx1, const int nx2, + const int nx3) { + BlockList_t block_list; + block_list.reserve(NBLOCKS); + for (int i = 0; i < NBLOCKS; ++i) { + auto pmb = std::make_shared(); + // Directly set cellbounds since it's public + if (nx3 > 0) { + pmb->cellbounds = parthenon::IndexShape(nx3, nx2, nx1, parthenon::Globals::nghost); + } else if (nx2 > 0) { + pmb->cellbounds = parthenon::IndexShape(nx2, nx1, parthenon::Globals::nghost); + } else { + pmb->cellbounds = parthenon::IndexShape(nx1, parthenon::Globals::nghost); + } + auto &pmbd = pmb->meshblock_data.Get(); + pmbd->Initialize(pkg, pmb); + block_list.push_back(pmb); + } + return block_list; +} + +// Test configuration struct +struct TestConfig { + int ndim; + int nx1, nx2, nx3; + int nghost; + int nkp, njp; + std::string description; + + int get_nk() const { return (ndim >= 3) ? nx3 : 1; } + int get_nj() const { return (ndim >= 2) ? nx2 : 1; } + int get_ni() const { return nx1; } +}; + } // namespace TEST_CASE("IndexSplit", "[IndexSplit]") { @@ -280,3 +319,456 @@ TEST_CASE("IndexSplit", "[IndexSplit]") { } } } + +TEST_CASE("IndexSplit Comprehensive", "[IndexSplit][comprehensive]") { + // Save original nghost value + const int original_nghost = parthenon::Globals::nghost; + + // Define test configurations focusing on cases with ghosts (the interesting cases!) + // Most tests use nghost=2 (typical for real simulations) + // Key: test memory layout with j-fusion (no_outer or small njp) + std::vector configs = { + // 1D cases - nghost matters less but include for completeness + {1, 4, 0, 0, 2, IndexSplit::all_outer, IndexSplit::no_outer, + "1D small ng=2 all_outer"}, + {1, 4, 0, 0, 2, 1, 1, "1D small ng=2 nkp=1"}, + {1, 16, 0, 0, 2, 4, 1, "1D medium ng=2 nkp=4 (divides evenly)"}, + {1, 16, 0, 0, 2, 5, 1, "1D medium ng=2 nkp=5 (doesn't divide)"}, + {1, 16, 0, 0, 3, 4, 1, "1D medium ng=3 nkp=4"}, + + // 2D cases - j-fusion starts to matter + {2, 4, 4, 0, 2, IndexSplit::all_outer, IndexSplit::no_outer, + "2D ng=2 all_outer,no_outer (full j-fusion)"}, + {2, 4, 4, 0, 2, IndexSplit::all_outer, IndexSplit::all_outer, + "2D ng=2 all_outer,all_outer (no j-fusion)"}, + {2, 4, 4, 0, 2, IndexSplit::no_outer, IndexSplit::no_outer, + "2D ng=2 no_outer,no_outer (all fused)"}, + {2, 6, 6, 0, 2, 3, 1, "2D ng=2 nkp=3 njp=1 (full j-fusion)"}, + {2, 6, 6, 0, 2, 4, 1, "2D ng=2 nkp=4 njp=1 (doesn't divide, full j-fusion)"}, + {2, 6, 6, 0, 2, 1, 3, "2D ng=2 njp=3 (j split, divides evenly)"}, + {2, 6, 6, 0, 2, 1, 4, "2D ng=2 njp=4 (j split, doesn't divide)"}, + {2, 6, 6, 0, 2, 3, 2, "2D ng=2 nkp=3 njp=2 (partial j-fusion)"}, + {2, 8, 8, 0, 3, 2, 2, "2D ng=3 nkp=2 njp=2"}, + + // 3D cases - the most important for IndexSplit + {3, 4, 4, 4, 2, IndexSplit::all_outer, IndexSplit::no_outer, + "3D ng=2 all_outer,no_outer (full j-fusion)"}, + {3, 4, 4, 4, 2, IndexSplit::all_outer, IndexSplit::all_outer, + "3D ng=2 all_outer,all_outer (no j-fusion)"}, + {3, 4, 4, 4, 2, IndexSplit::no_outer, IndexSplit::no_outer, + "3D ng=2 no_outer,no_outer (everything fused)"}, + {3, 4, 4, 4, 2, 2, 1, "3D ng=2 nkp=2 njp=1 (full j-fusion)"}, + {3, 6, 6, 6, 2, 3, 1, "3D ng=2 nkp=3 njp=1 (divides evenly, full j-fusion)"}, + {3, 6, 6, 6, 2, 4, 1, "3D ng=2 nkp=4 njp=1 (doesn't divide, full j-fusion)"}, + {3, 6, 6, 6, 2, 1, 3, "3D ng=2 njp=3 (j split, divides evenly)"}, + {3, 6, 6, 6, 2, 1, 4, "3D ng=2 njp=4 (j split, doesn't divide)"}, + {3, 6, 6, 6, 2, 3, 2, "3D ng=2 nkp=3 njp=2 (partial j-fusion)"}, + {3, 6, 6, 6, 2, 2, 2, "3D ng=2 nkp=2 njp=2 (partial j-fusion)"}, + {3, 8, 8, 8, 3, 2, 2, "3D ng=3 nkp=2 njp=2"}, + {3, 8, 8, 8, 4, 2, 1, "3D ng=4 nkp=2 njp=1 (full j-fusion, large ghosts)"}, + + // Asymmetric 3D case with ghosts + {3, 4, 8, 16, 2, 4, 2, "3D ng=2 asymmetric 4x8x16 njp=2"}, + {3, 4, 8, 16, 2, 4, 1, "3D ng=2 asymmetric 4x8x16 njp=1 (full j-fusion)"}, + + // Sanity check: one case with nghost=0 to verify it still works + {3, 4, 4, 4, 0, IndexSplit::no_outer, IndexSplit::no_outer, + "3D ng=0 no_outer,no_outer (sanity check)"}, + }; + + // Setup package for all tests + const std::vector scalar_shape{16, 16, 16}; + Metadata m({Metadata::Independent, Metadata::WithFluxes}, scalar_shape); + auto pkg = std::make_shared("Test package"); + pkg->AddField(v1::name(), m); + + for (const auto &config : configs) { + GIVEN(config.description) { + // Set nghost for this test + parthenon::Globals::nghost = config.nghost; + + // Create mesh blocks + constexpr int NBLOCKS = 3; + BlockList_t block_list; + if (config.nx2 == config.nx1 && config.nx3 == config.nx1) { + // Symmetric case - use simpler helper + block_list = MakeBlockList(pkg, NBLOCKS, config.nx1, config.ndim); + } else { + // Asymmetric case - use specialized helper + block_list = + MakeBlockListAsymmetric(pkg, NBLOCKS, config.nx1, config.nx2, config.nx3); + } + + MeshData mesh_data("base"); + mesh_data.Initialize(block_list, nullptr); + + WHEN("Using IndexDomain constructor") { + IndexSplit sp(&mesh_data, IndexDomain::interior, config.nkp, config.njp); + + // Get expected bounds for verification + auto kb = mesh_data.GetBoundsK(IndexDomain::interior); + auto jb = mesh_data.GetBoundsJ(IndexDomain::interior); + auto ib = mesh_data.GetBoundsI(IndexDomain::interior); + + THEN("outer_size() is correct") { + // Compute expected outer_size based on resolved nkp and njp + int expected_nkp = config.nkp; + int expected_njp = config.njp; + const int total_k = kb.e - kb.s + 1; + const int total_j = jb.e - jb.s + 1; + + if (expected_nkp == IndexSplit::all_outer) { + expected_nkp = total_k; + } else if (expected_nkp == IndexSplit::no_outer) { + expected_nkp = 1; + } else if (expected_nkp == 0) { +#ifdef PARTHENON_ENABLE_GPU + expected_nkp = total_k; +#else + expected_nkp = 1; +#endif + } + expected_nkp = std::min(expected_nkp, total_k); + + if (expected_njp == IndexSplit::all_outer) { + expected_njp = total_j; + } else if (expected_njp == IndexSplit::no_outer) { + expected_njp = 1; + } else if (expected_njp == 0) { +#ifdef PARTHENON_ENABLE_GPU + expected_njp = total_j; // Simplified - actual code is more complex +#else + expected_njp = 1; +#endif + } + expected_njp = std::min(expected_njp, total_j); + + REQUIRE(sp.outer_size() == expected_nkp * expected_njp); + } + + THEN("GetBounds methods cover the domain correctly") { + // Verify (k, j) pairs are covered exactly once + const int nk = kb.e - kb.s + 1; + const int nj = jb.e - jb.s + 1; + std::vector> kj_coverage(nk, std::vector(nj, 0)); + + for (int p = 0; p < sp.outer_size(); ++p) { + auto krange = sp.GetBoundsK(p); + auto jrange = sp.GetBoundsJ(p); + + // Verify ranges are within bounds + REQUIRE(krange.s >= kb.s); + REQUIRE(krange.e <= kb.e); + REQUIRE(jrange.s >= jb.s); + REQUIRE(jrange.e <= jb.e); + + // Mark all (k,j) pairs for this chunk + for (int k = krange.s; k <= krange.e; ++k) { + for (int j = jrange.s; j <= jrange.e; ++j) { + int kidx = k - kb.s; + int jidx = j - jb.s; + kj_coverage[kidx][jidx]++; + } + } + } + + // Verify every (k,j) pair is covered exactly once + for (int kidx = 0; kidx < nk; ++kidx) { + for (int jidx = 0; jidx < nj; ++jidx) { + REQUIRE(kj_coverage[kidx][jidx] == 1); + } + } + + // Verify i-range is consistent + auto ib_check = sp.GetBoundsI(0); + REQUIRE(ib_check.s == ib.s); + REQUIRE(ib_check.e == ib.e); + } + + THEN("get_i and get_deltaj decode inner indices correctly") { + auto kb_entire = mesh_data.GetBoundsK(IndexDomain::entire); + auto jb_entire = mesh_data.GetBoundsJ(IndexDomain::entire); + auto ib_entire = mesh_data.GetBoundsI(IndexDomain::entire); + const int ni_entire = ib_entire.e - ib_entire.s + 1; + + // Test for first outer index + if (sp.outer_size() > 0) { + auto jrange = sp.GetBoundsJ(0); + auto inner = sp.GetInnerBounds(jrange); + + // Test at start, middle, and end of inner range + std::vector test_indices = {inner.s}; + if (inner.e > inner.s) { + test_indices.push_back((inner.s + inner.e) / 2); + test_indices.push_back(inner.e); + } + + for (int idx : test_indices) { + int i = sp.get_i(idx); + int deltaj = sp.get_deltaj(idx); + + // Verify i is in valid range + REQUIRE(i >= 0); + REQUIRE(i < ni_entire); + + // Verify deltaj is in valid range + int expected_max_deltaj = jrange.e - jrange.s + 1; + REQUIRE(deltaj >= 0); + REQUIRE(deltaj < expected_max_deltaj); + + // Verify round-trip: reconstruct idx from i and deltaj + int reconstructed_idx = deltaj * ni_entire + i; + REQUIRE(reconstructed_idx == idx); + } + } + } + + THEN("get_max_ni/nj/nk methods return reasonable values") { + auto ib_entire = mesh_data.GetBoundsI(IndexDomain::entire); + auto jb_entire = mesh_data.GetBoundsJ(IndexDomain::entire); + auto kb_entire = mesh_data.GetBoundsK(IndexDomain::entire); + + REQUIRE(sp.get_max_ni() == ib_entire.e - ib_entire.s + 1); + REQUIRE(sp.get_max_nj() > 0); + REQUIRE(sp.get_max_nk() > 0); + REQUIRE(sp.get_max_nij() == sp.get_max_ni() * sp.get_max_nj()); + } + + THEN("Coverage test: every logical point visited exactly once") { + // Use the proper triple-nested loop structure + // Need bounds for ENTIRE domain since get_i/get_deltaj return coordinates + // in entire domain (includes ghosts that may be visited for memory contiguity) + auto kb_entire = mesh_data.GetBoundsK(IndexDomain::entire); + auto jb_entire = mesh_data.GetBoundsJ(IndexDomain::entire); + auto ib_entire = mesh_data.GetBoundsI(IndexDomain::entire); + const int nk_entire = kb_entire.e - kb_entire.s + 1; + const int nj_entire = jb_entire.e - jb_entire.s + 1; + const int ni_entire = ib_entire.e - ib_entire.s + 1; + + // Track coverage: how many times each (k,j,i) point is visited + // Must be sized to entire domain since inner loop can touch ghosts + using atomic_view = Kokkos::MemoryTraits; + Kokkos::View coverage("coverage", nk_entire, nj_entire, + ni_entire); + Kokkos::View counters("counters", 2); + // counters(0) = total_iterations + // counters(1) = ghost_iterations + + parthenon::par_for_outer( + DEFAULT_OUTER_LOOP_PATTERN, "Test IndexSplit Coverage", DevExecSpace(), 0, + 0, 0, sp.outer_size() - 1, + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int outer_idx) { + const auto krange = sp.GetBoundsK(outer_idx); + const auto jrange = sp.GetBoundsJ(outer_idx); + const auto inner = sp.GetInnerBounds(jrange); + + for (int k = krange.s; k <= krange.e; ++k) { + parthenon::par_for_inner(member, inner.s, inner.e, [&](const int idx) { + counters(0) += 1; // total iterations + + int i = sp.get_i(idx); + int deltaj = sp.get_deltaj(idx); + int j = jrange.s + deltaj; + + bool is_ghost = sp.is_ghost(outer_idx, k, idx); + if (is_ghost) { + counters(1) += 1; // ghost iterations + } + + // Track coverage - i,j,k are in entire domain coordinates + coverage(k - kb_entire.s, j - jb_entire.s, i - ib_entire.s) += 1; + }); + } + }); + + // Copy to host and verify + auto coverage_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), coverage); + auto counters_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), counters); + + int total_iterations = counters_h(0); + int ghost_iterations = counters_h(1); + int interior_iterations = total_iterations - ghost_iterations; + + const int nk = kb.e - kb.s + 1; + const int nj = jb.e - jb.s + 1; + const int ni = ib.e - ib.s + 1; + int expected_interior = nk * nj * ni; + + // Verify every INTERIOR point visited exactly once + // (Ghost points may be visited 0, 1, or more times - we don't care) + int interior_coverage_sum = 0; + for (int k = kb.s; k <= kb.e; ++k) { + for (int j = jb.s; j <= jb.e; ++j) { + for (int i = ib.s; i <= ib.e; ++i) { + int cov = coverage_h(k - kb_entire.s, j - jb_entire.s, i - ib_entire.s); + REQUIRE(cov == 1); + interior_coverage_sum += cov; + } + } + } + + // Verify counts + REQUIRE(interior_coverage_sum == expected_interior); + REQUIRE(interior_iterations == expected_interior); + REQUIRE(total_iterations == interior_iterations + ghost_iterations); + } + } + + WHEN("Using explicit IndexRange constructor") { + auto kb = mesh_data.GetBoundsK(IndexDomain::interior); + auto jb = mesh_data.GetBoundsJ(IndexDomain::interior); + auto ib = mesh_data.GetBoundsI(IndexDomain::interior); + + IndexSplit sp_explicit(&mesh_data, kb, jb, ib, config.nkp, config.njp); + IndexSplit sp_domain(&mesh_data, IndexDomain::interior, config.nkp, config.njp); + + THEN("Results match IndexDomain constructor") { + REQUIRE(sp_explicit.outer_size() == sp_domain.outer_size()); + + for (int p = 0; p < sp_explicit.outer_size(); ++p) { + auto kb_exp = sp_explicit.GetBoundsK(p); + auto kb_dom = sp_domain.GetBoundsK(p); + REQUIRE(kb_exp.s == kb_dom.s); + REQUIRE(kb_exp.e == kb_dom.e); + + auto jb_exp = sp_explicit.GetBoundsJ(p); + auto jb_dom = sp_domain.GetBoundsJ(p); + REQUIRE(jb_exp.s == jb_dom.s); + REQUIRE(jb_exp.e == jb_dom.e); + + auto ib_exp = sp_explicit.GetBoundsI(p); + auto ib_dom = sp_domain.GetBoundsI(p); + REQUIRE(ib_exp.s == ib_dom.s); + REQUIRE(ib_exp.e == ib_dom.e); + } + + REQUIRE(sp_explicit.get_max_ni() == sp_domain.get_max_ni()); + REQUIRE(sp_explicit.get_max_nj() == sp_domain.get_max_nj()); + REQUIRE(sp_explicit.get_max_nk() == sp_domain.get_max_nk()); + } + } + } + } + + // Restore original nghost + parthenon::Globals::nghost = original_nghost; +} + +TEST_CASE("IndexSplit Gold File Regression", "[IndexSplit][gold]") { + // Save original nghost value + const int original_nghost = parthenon::Globals::nghost; + + // Representative configurations to lock down behavior + std::vector gold_configs = { + {3, 4, 4, 4, 2, IndexSplit::all_outer, IndexSplit::no_outer, "3D_ng2_allk_noj"}, + {3, 4, 4, 4, 2, IndexSplit::all_outer, IndexSplit::all_outer, "3D_ng2_allk_allj"}, + {3, 4, 4, 4, 2, 2, 1, "3D_ng2_nkp2_njp1"}, + {3, 6, 6, 6, 2, 3, 2, "3D_ng2_nkp3_njp2"}, + {3, 4, 8, 16, 2, 4, 2, "3D_ng2_asym_4x8x16"}, + {2, 6, 6, 0, 2, IndexSplit::all_outer, 3, "2D_ng2_allk_njp3"}, + }; + + // Setup package + const std::vector scalar_shape{16, 16, 16}; + Metadata m({Metadata::Independent, Metadata::WithFluxes}, scalar_shape); + auto pkg = std::make_shared("Test package"); + pkg->AddField(v1::name(), m); + + for (const auto &config : gold_configs) { + GIVEN(config.description) { + parthenon::Globals::nghost = config.nghost; + + // Create mesh blocks + constexpr int NBLOCKS = 2; + BlockList_t block_list; + if (config.nx2 == config.nx1 && config.nx3 == config.nx1) { + block_list = MakeBlockList(pkg, NBLOCKS, config.nx1, config.ndim); + } else { + block_list = + MakeBlockListAsymmetric(pkg, NBLOCKS, config.nx1, config.nx2, config.nx3); + } + + MeshData mesh_data("base"); + mesh_data.Initialize(block_list, nullptr); + IndexSplit sp(&mesh_data, IndexDomain::interior, config.nkp, config.njp); + + // Get bounds for reference + auto kb = mesh_data.GetBoundsK(IndexDomain::interior); + auto jb = mesh_data.GetBoundsJ(IndexDomain::interior); + auto ib = mesh_data.GetBoundsI(IndexDomain::interior); + + WHEN("Recording structure to gold file") { + // Build gold file content + std::ostringstream gold_content; + gold_content << "# IndexSplit Gold File: " << config.description << "\n"; + gold_content << "# Config: ndim=" << config.ndim << " nx=" << config.nx1 << "," + << config.nx2 << "," << config.nx3 << " nghost=" << config.nghost + << " nkp=" << config.nkp << " njp=" << config.njp << "\n"; + gold_content << "# Domain bounds: k=[" << kb.s << "," << kb.e << "] " + << "j=[" << jb.s << "," << jb.e << "] " + << "i=[" << ib.s << "," << ib.e << "]\n"; + gold_content << "outer_size=" << sp.outer_size() << "\n"; + + for (int p = 0; p < sp.outer_size(); ++p) { + auto krange = sp.GetBoundsK(p); + auto jrange = sp.GetBoundsJ(p); + auto irange = sp.GetBoundsI(p); + auto inner = sp.GetInnerBounds(jrange); + int inner_size = inner.e - inner.s + 1; + + gold_content << "p=" << p << " k=[" << krange.s << "," << krange.e << "]" + << " j=[" << jrange.s << "," << jrange.e << "]" + << " i=[" << irange.s << "," << irange.e << "]" + << " inner=[" << inner.s << "," << inner.e << "]" + << " inner_size=" << inner_size << "\n"; + } + + std::string gold_str = gold_content.str(); + + // Try to read existing gold file + // Path is relative to build directory where tests run + std::string gold_path = + "../tst/unit/gold_files/index_split/" + config.description + ".gold"; + + // Check if we should generate gold files + const char *gen_gold = std::getenv("GENERATE_GOLD"); + if (gen_gold && std::string(gen_gold) == "1") { + // Write gold file + std::ofstream out_file(gold_path); + out_file << gold_str; + out_file.close(); + INFO("Generated gold file: " << gold_path); + REQUIRE(true); + } else { + std::ifstream gold_file(gold_path); + if (!gold_file.good()) { + // Gold file doesn't exist - skip test + INFO("Gold file does not exist: " << gold_path); + INFO("Run with GENERATE_GOLD=1 to create gold files"); + } else { + // Gold file exists - compare + std::stringstream existing_content; + existing_content << gold_file.rdbuf(); + std::string existing_str = existing_content.str(); + + if (gold_str != existing_str) { + // Mismatch - print both for debugging + INFO("Gold file mismatch for " << config.description); + INFO("Expected:\n" << existing_str); + INFO("Got:\n" << gold_str); + REQUIRE(gold_str == existing_str); + } else { + // Match - test passes + REQUIRE(true); + } + } + } + } + } + } + + // Restore original nghost + parthenon::Globals::nghost = original_nghost; +}