Fix MSA getData and accumulate operations (#3562)

Co-authored-by: Jaemin Choi <jchoi157@illinois.edu>
charmplusplus · Jan 27, 2022 · 5ed197e · 5ed197e
1 parent 2278b17
commit 5ed197e
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 6 deletions.
diff --git a/src/libs/ck-libs/multiphaseSharedArrays/msa-DistPageMgr.h b/src/libs/ck-libs/multiphaseSharedArrays/msa-DistPageMgr.h
@@ -393,8 +393,13 @@ class MSA_PageT {
 	/** Merger object */
 	MERGER m;
 
+	using self_type = MSA_PageT<ENTRY, MERGER, ENTRIES_PER_PAGE>;
+
 public:
 
+	static_assert(!std::is_same<ENTRY, bool>::value,
+				  "boolean MSAs currently unsupported");
+
 	MSA_PageT()
 		{
 			data.resize(ENTRIES_PER_PAGE);
@@ -429,8 +434,12 @@ class MSA_PageT {
 	// These accessors might be used by the templated code.
  	inline ENTRY &operator[](int i) {return data[i];}
  	inline const ENTRY &operator[](int i) const {return data[i];}
-    inline ENTRY *getData() { return data; }
-	inline const ENTRY *getData() const { return (const ENTRY*) data; }
+
+	inline ENTRY *getData() { return (ENTRY *)this->data.data(); }
+
+	inline const ENTRY *getData() const {
+		return const_cast<self_type *>(this)->getData();
+	}
 };
 
 //=============================== Cache Manager =================================
@@ -1270,7 +1279,7 @@ class MSA_PageArray : public CBase_MSA_PageArray<ENTRY_TYPE, ENTRY_OPS_CLASS, EN
     // MSA_PageArray::
     inline void combine(const ENTRY_TYPE* buffer, unsigned int begin, unsigned int end)
 		{
-			ENTRY_TYPE* pagePtr = epage.data() + begin;
+			ENTRY_TYPE* pagePtr = (ENTRY_TYPE *)epage.data() + begin;
 			for(unsigned int i = 0; i < (end - begin); i++)
 				entryOpsObject.accumulate(pagePtr[i], buffer[i]);
 		}
@@ -1307,7 +1316,7 @@ class MSA_PageArray : public CBase_MSA_PageArray<ENTRY_TYPE, ENTRY_OPS_CLASS, EN
 			if (entryOpsObject.pupEveryElement())
 				cache[pe].ReceivePageWithPUP(pageNo(), page_t(epage), epage.size());
 			else
-				cache[pe].ReceivePage(pageNo(), epage.data(), epage.size());
+				cache[pe].ReceivePage(pageNo(), (ENTRY_TYPE *)epage.data(), epage.size());
 		}
 
     /// Receive a non-runlength encoded page from the network:

diff --git a/src/libs/ck-libs/multiphaseSharedArrays/msa-distArray.h b/src/libs/ck-libs/multiphaseSharedArrays/msa-distArray.h
@@ -297,6 +297,13 @@ class MSAWrite : public MSAHandle<MSA>
     }
 };
 
+template<class MSA>
+class MSAFlatten;
+
+// type trait to map an n-dimension MSA to its 1D counterpart
+template<class MSA>
+using MSAFlattenT = typename MSAFlatten<MSA>::type;
+
 template<class MSA>
 class MSAAccum : public MSAHandle<MSA>
 {
@@ -316,7 +323,7 @@ class MSAAccum : public MSAHandle<MSA>
     inline Accumulable<MSA> accumulate(int x)
     {
         checkValid();
-        return Accumulable<MSA>(msa->accumulate(x));
+        return Accumulable<MSA>(((MSAFlattenT<MSA> *)msa)->accumulate(x));
     }
     inline Accumulable<MSA> operator() (int x)
     { return accumulate(x); }
@@ -408,7 +415,9 @@ class MSA1D
     friend class MSAHandle<thisMSA>;
     friend class MSARead<thisMSA>;
     friend class MSAWrite<thisMSA>;
-    friend class MSAAccum<thisMSA>;
+
+    template<typename MSA>
+    friend class MSAAccum;
 
 protected:
     /// Total number of ENTRY's in the whole array.
@@ -668,6 +677,10 @@ class MSA1D
     }
 };
 
+template<class ENTRY, class ENTRY_OPS_CLASS, unsigned int ENTRIES_PER_PAGE>
+struct MSAFlatten<MSA1D<ENTRY, ENTRY_OPS_CLASS, ENTRIES_PER_PAGE>> {
+    using type = MSA1D<ENTRY, ENTRY_OPS_CLASS, ENTRIES_PER_PAGE>;
+};
 
 // define a 2d distributed array based on the 1D array, support row major and column
 // major arrangement of data
@@ -792,6 +805,11 @@ class MSA2D : public MSA1D<ENTRY, ENTRY_OPS_CLASS, ENTRIES_PER_PAGE>
     }
 
 protected:
+    inline ENTRY& accumulate(unsigned int row, unsigned int col)
+    {
+        return super::accumulate(getIndex(row, col));
+    }
+
     inline const ENTRY& get(unsigned int row, unsigned int col)
     {
         return super::get(getIndex(row, col));
@@ -810,6 +828,11 @@ class MSA2D : public MSA1D<ENTRY, ENTRY_OPS_CLASS, ENTRIES_PER_PAGE>
     }
 };
 
+template<class ENTRY, class ENTRY_OPS_CLASS, unsigned int ENTRIES_PER_PAGE, MSA_Array_Layout_t ARRAY_LAYOUT>
+struct MSAFlatten<MSA2D<ENTRY, ENTRY_OPS_CLASS, ENTRIES_PER_PAGE, ARRAY_LAYOUT>> {
+    using type = MSA1D<ENTRY, ENTRY_OPS_CLASS, ENTRIES_PER_PAGE>;
+};
+
 /**
    The MSA3D class is a handle to a distributed shared array of items
    of data type ENTRY. There are nEntries total numer of ENTRY's, with
@@ -1149,5 +1172,10 @@ class MSA3D
     }
 };
 
+template<class ENTRY, class ENTRY_OPS_CLASS, unsigned int ENTRIES_PER_PAGE>
+struct MSAFlatten<MSA3D<ENTRY, ENTRY_OPS_CLASS, ENTRIES_PER_PAGE>> {
+    using type = MSA1D<ENTRY, ENTRY_OPS_CLASS, ENTRIES_PER_PAGE>;
+};
+
 }
 #endif