From 19a4cd2467aace7fb4b7b7f264c91d5218eef939 Mon Sep 17 00:00:00 2001
From: Christian Zimmermann <christian.zimmermann@ur.de>
Date: Wed, 16 Oct 2024 23:32:36 -0700
Subject: [PATCH] mpi: more on rank format (TODO: think about pos)

---
 src/opt/mpi/include/rrange.cc.h | 31 +++++++++++++++++++++++++------
 src/opt/mpi/include/rrange.h    | 12 ++++++++++--
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/src/opt/mpi/include/rrange.cc.h b/src/opt/mpi/include/rrange.cc.h
index b96e285..81e2e99 100644
--- a/src/opt/mpi/include/rrange.cc.h
+++ b/src/opt/mpi/include/rrange.cc.h
@@ -32,7 +32,8 @@ namespace CNORXZ
 	    mK(std::make_shared<IndexK>(mRange->geom())),
 	    mNRanks(getNumRanks()),
 	    mRankFormat(in.mRankFormat),
-	    mRankOffset(in.mRankOffset)
+	    mRankOffset(in.mRankOffset),
+	    mStepRatio(in.mStepRatio)
 	{
 	    *this = in.lex();
 	}
@@ -46,6 +47,7 @@ namespace CNORXZ
 	    mNRanks = getNumRanks();
 	    mRankOffset = in.mRankOffset;
 	    mRankFormat = in.mRankFormat;
+	    mStepRatio = mStepRatio;
 	    *this = in.lex();
 	    return *this;
 	}
@@ -57,7 +59,8 @@ namespace CNORXZ
 	    mK(std::make_shared<IndexK>(mRange->geom())),
 	    mNRanks(getNumRanks()),
 	    mRankFormat(1),
-	    mRankOffset(0)
+	    mRankOffset(0),
+	    mStepRatio(mI->pmax().val())
 	{
 	    *this = lexpos;
 	}
@@ -69,7 +72,8 @@ namespace CNORXZ
 	    mK(k),
 	    mNRanks(getNumRanks()),
 	    mRankFormat(1),
-	    mRankOffset(0)
+	    mRankOffset(0),
+	    mStepRatio(mI->pmax().val())
 	{
 	    (*this)();
 	}
@@ -133,6 +137,8 @@ namespace CNORXZ
 		*mK = pos / mI->lmax().val();
 		*mI = pos % mI->lmax().val();
 	    }
+	    // replace line once we know what to do !!!
+	    //IB::mPos = mK->pos() * mStepRatio + mI->pos();
 	    IB::mPos = mK->pos() * mI->pmax().val() + mI->pos();
 	    return *this;
 	}
@@ -399,6 +405,7 @@ namespace CNORXZ
 		    }
 		}
 	    }
+	    //IB::mPos = mK->pos() * mStepRatio + mI->pos();
 	    IB::mPos = mK->pos() * mI->pmax().val() + mI->pos();
 	    return *this;
 	}
@@ -422,15 +429,21 @@ namespace CNORXZ
 	}
 
 	template <class IndexI, class IndexK>
-	void RIndex<IndexI,IndexK>::setRankFormat(SizeT rankFormat)
+	void RIndex<IndexI,IndexK>::setRankFormat(SizeT rankFormat, SizeT stepRatio)
 	{
 	    const SizeT nr = getRankNumber();
 	    mRankFormat = rankFormat;
+	    CXZ_ASSERT(mNRanks % mRankFormat == 0, "rankFormat (" << mRankFormat
+		       << ") does not divide total number of ranks (" << mNRanks << ")");
+	    mStepRatio = stepRatio;
+	    if(mStepRatio == 0){
+		mStepRatio = mI->pmax().val();
+	    }
+	    CXZ_ASSERT(mStepRatio % mRankFormat == 0, "rankFormat (" << mRankFormat
+		       << ") does not divide stepRatio (" << mStepRatio << ")");
 	    const SizeT kmax = mK->pmax().val();
 	    const SizeT nextblocks = kmax * mRankFormat;
 	    mRankOffset = nr % mRankFormat + ( nr / nextblocks ) * nextblocks;
-	    CXZ_ASSERT(mNRanks % mRankFormat == 0, "rankFormat (" << mRankFormat
-		       << ") does not divide total number of ranks (" << mNRanks << ")");
 	}
 	
 	template <class IndexI, class IndexK>
@@ -439,6 +452,12 @@ namespace CNORXZ
 	    return mRankFormat;
 	}
 
+	template <class IndexI, class IndexK>
+	SizeT RIndex<IndexI,IndexK>::stepRatio() const
+	{
+	    return mStepRatio;
+	}
+
 	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 6adc66b..3f39fe8 100644
--- a/src/opt/mpi/include/rrange.h
+++ b/src/opt/mpi/include/rrange.h
@@ -159,11 +159,18 @@ namespace CNORXZ
 	    /** Get index indicating the current rank this index points to. */
 	    Sptr<IndexK> rankI() const;
 
-	    /** Set the rank format. */
-	    void setRankFormat(SizeT rankFormat);
+	    /** Set the rank format.
+		@param rankFormat Format of the rank index K.
+		@param stepRatio stepSize of K over stepsize of I; needs to be dividable by rankFormat.
+		If stepRatio == 0, it is set to the default value = mI->pmax().val().
+	     */
+	    void setRankFormat(SizeT rankFormat, SizeT stepRatio = 0);
 
 	    /** Get the rank format. */
 	    SizeT rankFormat() const;
+
+	    /** Get step ratio. */
+	    SizeT stepRatio() const;
 	    
 	private:
 	    SizeT mLex = 0;
@@ -173,6 +180,7 @@ namespace CNORXZ
 	    SizeT mNRanks; /**< Number of ranks; in general different from but dividable by mK's maximum. */
 	    SizeT mRankFormat = 1; /**< Frequency of ranks to be served by this index. */
 	    SizeT mRankOffset = 0; /** < Offset in case this index only serves a sub-slice of ranks .*/
+	    SizeT mStepRatio; /**< Ratio between the stepsizes of mK and mI. */
 	};
 
 	template <class IndexI, class IndexK>