I Rearranged the load balancer, so WSLB->NeighborLB, NeighborLB->NborBaseLB,
[charm.git] / src / ck-ldb / WSLB.h
index da2254ef4a831d389fc3e17d5e579a25aecc23aa..f0f5ce24b6372c17bb705737fab59f01aa109f9c 100644 (file)
@@ -1,18 +1,54 @@
-#ifndef _WSLB_H_
-#define _WSLB_H_
+#ifndef NEIGHBORLB_H
+#define NEIGHBORLB_H
 
-#include <math.h>
-
-#include "NeighborLB.h"
+#include <LBDatabase.h>
 #include "WSLB.decl.h"
 
 void CreateWSLB();
 
-class WSLB : public NeighborLB {
+class WSLBStatsMsg;
+class WSLBMigrateMsg;
+
+class WSLB : public Group
+{
 public:
   WSLB();
-private:
-  CmiBool QueryBalanceNow(int step) { return CmiTrue; };
+  ~WSLB();
+  static void staticAtSync(void*);
+  void AtSync(void); // Everything is at the PE barrier
+
+  void ReceiveStats(WSLBStatsMsg *);           // Receive stats on PE 0
+  void ResumeClients();
+  void ReceiveMigration(WSLBMigrateMsg *);     // Receive migration data
+
+  // Migrated-element callback
+  static void staticMigrated(void* me, LDObjHandle h);
+  void Migrated(LDObjHandle h);
+
+  void MigrationDone(void);  // Call when migration is complete
+  int step() { return mystep; };
+
+  struct MigrateInfo {  // Used in WSLBMigrateMsg
+    LDObjHandle obj;
+    int from_pe;
+    int to_pe;
+  };
+
+  struct LDStats {  // Passed to Strategy
+    int from_pe;
+    double total_walltime;
+    double total_cputime;
+    double idletime;
+    double bg_walltime;
+    double bg_cputime;
+    double obj_walltime;
+    double obj_cputime;
+    int proc_speed;
+  };
+
+protected:
+  virtual CmiBool QueryBalanceNow(int) { return CmiTrue; };  
+  virtual WSLBMigrateMsg* Strategy(LDStats* stats,int count);
   virtual int num_neighbors() {
     return (CkNumPes() > 5) ? 4 : (CkNumPes()-1);
   };
@@ -33,7 +69,64 @@ private:
       _n[3] = (me + npe - bigstep) % npe;
   };
 
-  NLBMigrateMsg* Strategy(NeighborLB::LDStats* stats, int count);
+  LBDatabase* theLbdb;
+  struct {
+    int proc_speed;
+    double total_walltime;
+    double total_cputime;
+    double idletime;
+    double bg_walltime;
+    double bg_cputime;
+    int obj_data_sz;
+    LDObjData* objData;
+    int comm_data_sz;
+    LDCommData* commData;
+    double obj_walltime;
+    double obj_cputime;
+  } myStats;
+
+private:
+  void FindNeighbors();
+  WSLBStatsMsg* AssembleStats();
+
+  int mystep;
+  int stats_msg_count;
+  WSLBStatsMsg** statsMsgsList;
+  LDStats* statsDataList;
+  int migrates_completed;
+  int migrates_expected;
+  WSLBMigrateMsg** mig_msgs;
+  int mig_msgs_received;
+  int mig_msgs_expected;
+  int* neighbor_pes;
+  int receive_stats_ready;
+  double start_lb_time;
 };
 
-#endif /* _WSLB_H_ */
+class WSLBStatsMsg : public CMessage_WSLBStatsMsg {
+public:
+  int from_pe;
+  int serial;
+  int proc_speed;
+  double total_walltime;
+  double total_cputime;
+  double idletime;
+  double bg_walltime;
+  double bg_cputime;
+  double obj_walltime;
+  double obj_cputime;
+}; 
+
+class WSLBMigrateMsg : public CMessage_WSLBMigrateMsg {
+public:
+  int n_moves;
+  WSLB::MigrateInfo* moves;
+
+  // Other methods & data members 
+  
+  static void* alloc(int msgnum, size_t size, int* array, int priobits); 
+  static void* pack(WSLBMigrateMsg* in); 
+  static WSLBMigrateMsg* unpack(void* in); 
+}; 
+
+#endif /* NEIGHBORLB_H */