From 03afacb1e5d3b86a7d3923e34320d9ca556959c2 Mon Sep 17 00:00:00 2001
From: Christian Zimmermann <christian.zimmermann@ur.de>
Date: Fri, 18 Oct 2024 23:59:41 -0700
Subject: [PATCH] mpi: rindex: synchronize/localize utilities

---
 src/opt/mpi/include/rrange.cc.h           | 62 +++++++++++++++++++++--
 src/opt/mpi/include/rrange.h              | 12 +++++
 src/opt/mpi/tests/roperation_unit_test.cc |  1 +
 3 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/src/opt/mpi/include/rrange.cc.h b/src/opt/mpi/include/rrange.cc.h
index 81e2e99..0a63253 100644
--- a/src/opt/mpi/include/rrange.cc.h
+++ b/src/opt/mpi/include/rrange.cc.h
@@ -339,10 +339,10 @@ namespace CNORXZ
 	template <class Xpr, class F>
 	constexpr decltype(auto) RIndex<IndexI,IndexK>::ifor(const Xpr& xpr, F&& f) const
 	{
-	    const SizeT r = getRankNumber();
-	    const SizeT off = r % mRankFormat;
-	    //assert(off == mRankOffset);
-	    const SizeT n = ( (r-off) / mRankFormat ) % mK->pmax().val();
+	    //const SizeT r = getRankNumber();
+	    //const SizeT off = r % mRankFormat;
+	    //const SizeT n = ( (r-off) / mRankFormat ) % mK->pmax().val();
+	    const SizeT n = mK->lex();
 	    return accxpr( n, mK->id(), mI->ifor(xpr, std::forward<F>(f)), NoF {});
 	}
 
@@ -458,6 +458,60 @@ namespace CNORXZ
 	    return mStepRatio;
 	}
 
+	template <class IndexI, class IndexK>
+	bool RIndex<IndexI,IndexK>::isSynchronous() const
+	{
+	    const size_t lx = lex();
+	    Vector<size_t> lexs(mNRanks);
+	    MPI_Allgather(&lx, 1, MPI_UNSIGNED_LONG, lexs.data(), 1, MPI_UNSIGNED_LONG,
+			  MPI_COMM_WORLD);
+	    const size_t glex = lexs[0];
+	    bool o = true;
+	    for(const auto& l: lexs){
+		if(l != glex){
+		    o = false;
+		    break;
+		}
+	    }
+	    return o;
+	}
+
+	template <class IndexI, class IndexK>
+	bool RIndex<IndexI,IndexK>::isLocalized() const
+	{
+	    const size_t lx = mK->lex();
+	    const unsigned short is = (lx == myrank());
+	    Vector<unsigned short> iss(mNRanks);
+	    MPI_Allgather(&is, 1, MPI_UNSIGNED_SHORT, iss.data(), 1, MPI_UNSIGNED_SHORT,
+			  MPI_COMM_WORLD);
+	    bool o = true;
+	    for(const auto& l: iss){
+		if(not l){
+		    o = false;
+		    break;
+		}
+	    }
+	    return o;
+	}
+	
+	template <class IndexI, class IndexK>
+	RIndex<IndexI,IndexK>& RIndex<IndexI,IndexK>::localize()
+	{
+	    (*mK) = myrank();
+	    (*this)();
+	    return *this;
+	}
+	
+	template <class IndexI, class IndexK>
+	SizeT RIndex<IndexI,IndexK>::myrank() const
+	{
+	    const SizeT r = getRankNumber();
+	    const SizeT off = r % mRankFormat;
+	    assert((r-off)/mRankFormat == r/mRankFormat);
+	    const SizeT n = ( (r-off) / mRankFormat ) % mK->pmax().val();
+	    return n;
+	}
+
 	template <class IndexI, class IndexK, class I1>
 	decltype(auto) operator*(const Sptr<RIndex<IndexI,IndexK>>& a, const Sptr<I1>& b)
 	{
diff --git a/src/opt/mpi/include/rrange.h b/src/opt/mpi/include/rrange.h
index 3f39fe8..c596f8b 100644
--- a/src/opt/mpi/include/rrange.h
+++ b/src/opt/mpi/include/rrange.h
@@ -171,6 +171,18 @@ namespace CNORXZ
 
 	    /** Get step ratio. */
 	    SizeT stepRatio() const;
+
+	    /** Check if index points to the same global position on each rank. */
+	    bool isSynchronous() const;
+
+	    /** Check if index points to own rank on each rank. */
+	    bool isLocalized() const;
+
+	    /** Set internl rank index to own rank. Results in non-synchronous state. */
+	    RIndex& localize();
+
+	    /** Get this rank considering rank format. */
+	    SizeT myrank() const;
 	    
 	private:
 	    SizeT mLex = 0;
diff --git a/src/opt/mpi/tests/roperation_unit_test.cc b/src/opt/mpi/tests/roperation_unit_test.cc
index 8c57ea0..eff2abc 100644
--- a/src/opt/mpi/tests/roperation_unit_test.cc
+++ b/src/opt/mpi/tests/roperation_unit_test.cc
@@ -140,6 +140,7 @@ namespace
 	for(const auto& r: *imap1){
 	    req[r] = true;
 	}
+	x->localize();
 	res.load(x, AB, req); // DUMMY, not used...
 	mM1.load(xp, AB, req);
 	res(x*A*B) = mapXpr(xp,x,imap1, mM1(xp*A*B) - mM1(x*A*B) );