Enumerator for an optimal federated plan at the program level

Author: min-guk

src/main/java/org/apache/sysds/hops/fedplanner/FederatedMemoTable.java

@@ -20,7 +20,6 @@
 package org.apache.sysds.hops.fedplanner;
 
 import org.apache.sysds.hops.Hop;
-import org.apache.sysds.hops.OptimizerUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.sysds.runtime.instructions.fed.FEDInstruction.FederatedOutput;
@@ -29,8 +28,6 @@
 import java.util.List;
 import java.util.ArrayList;
 import java.util.Map;
-import java.util.HashSet;
-import java.util.Set;
 
 /**
  * A Memoization Table for managing federated plans (FedPlan) based on combinations of Hops and fedOutTypes.
@@ -127,12 +124,12 @@ public void pruneFedPlan(long hopID, FederatedOutput federatedOutput) {
 	public static class HopCommon {
 		protected final Hop hopRef;         // Reference to the associated Hop
 		protected double selfCost;          // Current execution cost (compute + memory access)
-		protected double netTransferCost;   // Network transfer cost
+		protected double forwardingCost;   // Network transfer cost
 
 		protected HopCommon(Hop hopRef) {
 			this.hopRef = hopRef;
 			this.selfCost = 0;
-			this.netTransferCost = 0;
+			this.forwardingCost = 0;
 		}
 	}
 
@@ -175,7 +172,7 @@ public void prune() {
 	 * This class contains:
 	 * 1. selfCost: Cost of current hop (compute + input/output memory access)
 	 * 2. totalCost: Cumulative cost including this plan and all child plans
-	 * 3. netTransferCost: Network transfer cost for this plan to parent plan.
+	 * 3. forwardingCost: Network transfer cost for this plan to parent plan.
 	 * 
 	 * FedPlan is linked to FedPlanVariants, which in turn uses HopCommon to manage common properties and costs.
 	 */
@@ -192,14 +189,15 @@ public FedPlan(List<Pair<Long, FederatedOutput>> childFedPlans, FedPlanVariants
 
 		public void setTotalCost(double totalCost) {this.totalCost = totalCost;}
 		public void setSelfCost(double selfCost) {fedPlanVariants.hopCommon.selfCost = selfCost;}
-		public void setNetTransferCost(double netTransferCost) {fedPlanVariants.hopCommon.netTransferCost = netTransferCost;}
-		
+		public void setForwardingCost(double forwardingCost) {fedPlanVariants.hopCommon.forwardingCost = forwardingCost;}
+		public void applyIterationWeight(int iteration) {totalCost *= iteration;}
+
 		public Hop getHopRef() {return fedPlanVariants.hopCommon.hopRef;}
 		public long getHopID() {return fedPlanVariants.hopCommon.hopRef.getHopID();}
 		public FederatedOutput getFedOutType() {return fedPlanVariants.fedOutType;}
 		public double getTotalCost() {return totalCost;}
 		public double getSelfCost() {return fedPlanVariants.hopCommon.selfCost;}
-		public double getNetTransferCost() {return fedPlanVariants.hopCommon.netTransferCost;}
+		public double setForwardingCost() {return fedPlanVariants.hopCommon.forwardingCost;}
 		public List<Pair<Long, FederatedOutput>> getChildFedPlans() {return childFedPlans;}
 
 		/**
@@ -208,9 +206,9 @@ public FedPlan(List<Pair<Long, FederatedOutput>> childFedPlans, FedPlanVariants
 		 * @param parentFedOutType The federated output type of the parent plan.
 		 * @return The conditional network transfer cost.
 		 */
-		public double getCondNetTransferCost(FederatedOutput parentFedOutType) {
+		public double getCondForwardingCost(FederatedOutput parentFedOutType) {
 			if (parentFedOutType == getFedOutType()) return 0;
-			return fedPlanVariants.hopCommon.netTransferCost;
+			return fedPlanVariants.hopCommon.forwardingCost;
 		}
 	}
 }

src/main/java/org/apache/sysds/hops/fedplanner/FederatedMemoTablePrinter.java

@@ -67,7 +67,7 @@ private static void printFedPlanTreeRecursive(FederatedMemoTable.FedPlan plan, F
         sb.append(String.format(" {Total: %.1f, Self: %.1f, Net: %.1f}",
                 plan.getTotalCost(),
                 plan.getSelfCost(),
-                plan.getNetTransferCost()));
+                plan.setForwardingCost()));
 
         // Add matrix characteristics
         sb.append(" [")

src/main/java/org/apache/sysds/hops/fedplanner/FederatedPlanCostEnumerator.java

@@ -31,6 +31,16 @@
 import org.apache.sysds.hops.Hop;
 import org.apache.sysds.hops.fedplanner.FederatedMemoTable.FedPlan;
 import org.apache.sysds.hops.fedplanner.FederatedMemoTable.FedPlanVariants;
+import org.apache.sysds.parser.DMLProgram;
+import org.apache.sysds.parser.ForStatement;
+import org.apache.sysds.parser.ForStatementBlock;
+import org.apache.sysds.parser.FunctionStatement;
+import org.apache.sysds.parser.FunctionStatementBlock;
+import org.apache.sysds.parser.IfStatement;
+import org.apache.sysds.parser.IfStatementBlock;
+import org.apache.sysds.parser.StatementBlock;
+import org.apache.sysds.parser.WhileStatement;
+import org.apache.sysds.parser.WhileStatementBlock;
 import org.apache.sysds.runtime.instructions.fed.FEDInstruction.FederatedOutput;
 
 /**
@@ -39,16 +49,91 @@
  * to compute their costs.
  */
 public class FederatedPlanCostEnumerator {
+	public static void enumerateProgram(DMLProgram prog) {
+		for(StatementBlock sb : prog.getStatementBlocks())
+			enumerateStatementBlock(sb);
+	}
+
+	/**
+	 * Recursively enumerates federated execution plans for a given statement block.
+	 * This method processes each type of statement block (If, For, While, Function, and generic)
+	 * to determine the optimal federated plan.
+	 * 
+	 * @param sb The statement block to enumerate.
+	 */
+	public static void enumerateStatementBlock(StatementBlock sb) {
+		// While enumerating the program, recursively determine the optimal FedPlan and MemoTable
+		// for each statement block and statement.
+		// 1. How to recursively integrate optimal FedPlans and MemoTables across statements and statement blocks?
+			// 1) Is it determined using the same dynamic programming approach, or simply by summing the minimal plans?
+		// 2. Is there a need to share the MemoTable? Are there data/hop dependencies between statements?
+		// 3. How to predict the number of iterations for For and While loops?
+			// 1) If from/to/increment are constants: Calculations can be done at compile time.
+			// 2) If they are variables: Use default values at compile time, adjust at runtime, or predict using ML models.
+
+		if (sb instanceof IfStatementBlock) {
+			IfStatementBlock isb = (IfStatementBlock) sb;
+			IfStatement istmt = (IfStatement)isb.getStatement(0);
+
+			enumerateFederatedPlanCost(isb.getPredicateHops());
+
+			for (StatementBlock csb : istmt.getIfBody())
+				enumerateStatementBlock(csb);
+			for (StatementBlock csb : istmt.getElseBody())
+				enumerateStatementBlock(csb);
+
+			// Todo: 1. apply iteration weight to csbFedPlans (if: 0.5, else: 0.5)
+			// Todo: 2. Merge predFedPlans
+		} else if (sb instanceof ForStatementBlock) { //incl parfor
+			ForStatementBlock fsb = (ForStatementBlock) sb;
+
+			ForStatement fstmt = (ForStatement)fsb.getStatement(0);
+
+			enumerateFederatedPlanCost(fsb.getFromHops());
+			enumerateFederatedPlanCost(fsb.getToHops());
+			enumerateFederatedPlanCost(fsb.getIncrementHops());
+
+			for (StatementBlock csb : fstmt.getBody())
+				enumerateStatementBlock(csb);
+
+			// Todo: 1. get(predict) # of Iterations
+			// Todo: 2. apply iteration weight to csbFedPlans
+			// Todo: 3. Merge csbFedPlans and predFedPlans
+		} else if (sb instanceof WhileStatementBlock) {
+			WhileStatementBlock wsb = (WhileStatementBlock) sb;
+			WhileStatement wstmt = (WhileStatement)wsb.getStatement(0);
+			enumerateFederatedPlanCost(wsb.getPredicateHops());
+
+			ArrayList<FedPlan> csbFedPlans = new ArrayList<>();
+			for (StatementBlock csb : wstmt.getBody())
+				enumerateStatementBlock(csb);
+
+			// Todo: 1. get(predict) # of Iterations
+			// Todo: 2. apply iteration weight to csbFedPlans
+			// Todo: 3. Merge csbFedPlans and predFedPlans
+		} else  if (sb instanceof FunctionStatementBlock) {
+			FunctionStatementBlock fsb = (FunctionStatementBlock)sb;
+			FunctionStatement fstmt = (FunctionStatement)fsb.getStatement(0);
+			for (StatementBlock csb : fstmt.getBody())
+				enumerateStatementBlock(csb);
+
+			// Todo: 1. Merge csbFedPlans
+		} else { //generic (last-level)
+			if( sb.getHops() != null )
+				for( Hop c : sb.getHops() )
+					enumerateFederatedPlanCost(c);
+		}
+	}
+
 	/**
 	 * Entry point for federated plan enumeration. This method creates a memo table
 	 * and returns the minimum cost plan for the entire Directed Acyclic Graph (DAG).
 	 * It also resolves conflicts where FedPlans have different FederatedOutput types.
 	 * 
 	 * @param rootHop The root Hop node from which to start the plan enumeration.
-	 * @param printTree A boolean flag indicating whether to print the federated plan tree.
 	 * @return The optimal FedPlan with the minimum cost for the entire DAG.
 	 */
-	public static FedPlan enumerateFederatedPlanCost(Hop rootHop, boolean printTree) {
+	public static FedPlan enumerateFederatedPlanCost(Hop rootHop) {
 		// Create new memo table to store all plan variants
 		FederatedMemoTable memoTable = new FederatedMemoTable();
 
@@ -61,8 +146,8 @@ public static FedPlan enumerateFederatedPlanCost(Hop rootHop, boolean printTree)
 		// Detect conflicts in the federated plans where different FedPlans have different FederatedOutput types
 		double additionalTotalCost = detectAndResolveConflictFedPlan(optimalPlan, memoTable);
 
-		// Optionally print the federated plan tree if requested
-		if (printTree) FederatedMemoTablePrinter.printFedPlanTree(optimalPlan, memoTable, additionalTotalCost);
+		// Print the federated plan tree if requested
+		FederatedMemoTablePrinter.printFedPlanTree(optimalPlan, memoTable, additionalTotalCost);
 
 		return optimalPlan;
 	}

src/main/java/org/apache/sysds/hops/fedplanner/FederatedPlanCostEstimator.java

@@ -61,7 +61,7 @@ public static void computeFederatedPlanCost(FedPlan currentPlan, FederatedMemoTa
 			totalCost = computeCurrentCost(currentHop);
 			currentPlan.setSelfCost(totalCost);
 			// Calculate potential network transfer cost if federation type changes
-			currentPlan.setNetTransferCost(computeHopNetworkAccessCost(currentHop.getOutputMemEstimate()));
+			currentPlan.setForwardingCost(computeHopNetworkAccessCost(currentHop.getOutputMemEstimate()));
 		} else {
 			totalCost = currentPlan.getSelfCost();
 		}
@@ -74,7 +74,7 @@ public static void computeFederatedPlanCost(FedPlan currentPlan, FederatedMemoTa
 			FedPlan planRef = memoTable.getMinCostFedPlan(childPlanPair);
 
 			// Add child plan cost (includes network transfer cost if federation types differ)
-			totalCost += planRef.getTotalCost() + planRef.getCondNetTransferCost(currentPlan.getFedOutType());
+			totalCost += planRef.getTotalCost() + planRef.getCondForwardingCost(currentPlan.getFedOutType());
 		}
 
 		// Step 3: Set final cumulative cost including current node
@@ -111,8 +111,8 @@ public static LinkedHashMap<FedPlan, Boolean> resolveConflictFedPlan(FederatedMe
 
 			// Flags to check if the plan involves network transfer
 			// Network transfer cost is calculated only once, even if it occurs multiple times
-			boolean isLOutNetTransfer = false;
-			boolean isFOutNetTransfer = false; 
+			boolean isLOutForwarding = false;
+			boolean isFOutForwarding = false;
 
 			// Determine the optimal federated output type based on the calculated costs
 			FederatedOutput optimalFedOutType;
@@ -143,35 +143,35 @@ public static LinkedHashMap<FedPlan, Boolean> resolveConflictFedPlan(FederatedMe
 					if (conflictParentFedPlan.getFedOutType() == FederatedOutput.LOUT) {
 						// (CASE 1) Previously, calculated was LOUT and parent was LOUT, so no network transfer cost occurred
 						// (CASE 5) If changing from calculated LOUT to current FOUT, network transfer cost occurs, but calculated later
-						isFOutNetTransfer = true;
+						isFOutForwarding = true;
 					} else {
 						// Previously, calculated was LOUT and parent was FOUT, so network transfer cost occurred
                     	// (CASE 2) If maintaining calculated LOUT to current LOUT, subtract existing network transfer cost and calculate later
-						isLOutNetTransfer = true;
-						lOutAdditionalCost -= confilctLOutFedPlan.getNetTransferCost();
+						isLOutForwarding = true;
+						lOutAdditionalCost -= confilctLOutFedPlan.setForwardingCost();
 
 						// (CASE 6) If changing from calculated LOUT to current FOUT, no network transfer cost occurs, so subtract it
-						fOutAdditionalCost -= confilctLOutFedPlan.getNetTransferCost();
+						fOutAdditionalCost -= confilctLOutFedPlan.setForwardingCost();
 					}
 				} else {
 					lOutAdditionalCost += confilctLOutFedPlan.getTotalCost() - confilctFOutFedPlan.getTotalCost();
 
 					if (conflictParentFedPlan.getFedOutType() == FederatedOutput.FOUT) {
-						isLOutNetTransfer = true;
+						isLOutForwarding = true;
 					} else {
-						isFOutNetTransfer = true;
-						lOutAdditionalCost -= confilctLOutFedPlan.getNetTransferCost();
-						fOutAdditionalCost -= confilctLOutFedPlan.getNetTransferCost();
+						isFOutForwarding = true;
+						lOutAdditionalCost -= confilctLOutFedPlan.setForwardingCost();
+						fOutAdditionalCost -= confilctLOutFedPlan.setForwardingCost();
 					}
 				}
 			}
 
 			// Add network transfer costs if applicable
-			if (isLOutNetTransfer) {
-				lOutAdditionalCost += confilctLOutFedPlan.getNetTransferCost();
+			if (isLOutForwarding) {
+				lOutAdditionalCost += confilctLOutFedPlan.setForwardingCost();
 			}
-			if (isFOutNetTransfer) {
-				fOutAdditionalCost += confilctFOutFedPlan.getNetTransferCost();
+			if (isFOutForwarding) {
+				fOutAdditionalCost += confilctFOutFedPlan.setForwardingCost();
 			}
 
 			// Determine the optimal federated output type based on the calculated costs

src/test/java/org/apache/sysds/test/component/federated/FederatedPlanCostEnumeratorTest.java

@@ -45,7 +45,7 @@ public class FederatedPlanCostEnumeratorTest extends AutomatedTestBase
 
 	@Override
 	public void setUp() {}
-	
+
 	@Test
 	public void testFederatedPlanCostEnumerator1() { runTest("FederatedPlanCostEnumeratorTest1.dml"); }
 
@@ -55,6 +55,21 @@ public void setUp() {}
 	@Test
 	public void testFederatedPlanCostEnumerator3() { runTest("FederatedPlanCostEnumeratorTest3.dml"); }
 
+	@Test
+	public void testFederatedPlanCostEnumerator4() { runTest("FederatedPlanCostEnumeratorTest4.dml"); }
+
+	@Test
+	public void testFederatedPlanCostEnumerator5() { runTest("FederatedPlanCostEnumeratorTest5.dml"); }
+
+	@Test
+	public void testFederatedPlanCostEnumerator6() { runTest("FederatedPlanCostEnumeratorTest6.dml"); }
+
+	@Test
+	public void testFederatedPlanCostEnumerator7() { runTest("FederatedPlanCostEnumeratorTest7.dml"); }
+
+	@Test
+	public void testFederatedPlanCostEnumerator8() { runTest("FederatedPlanCostEnumeratorTest4.dml"); }
+
 	// Todo: Need to write test scripts for the federated version
 	private void runTest( String scriptFilename ) {
 		int index = scriptFilename.lastIndexOf(".dml");
@@ -80,8 +95,7 @@ private void runTest( String scriptFilename ) {
 			dmlt.rewriteHopsDAG(prog);
 			dmlt.constructLops(prog);
 
-			Hop hops = prog.getStatementBlocks().get(0).getHops().get(0);
-			FederatedPlanCostEnumerator.enumerateFederatedPlanCost(hops, true);
+			FederatedPlanCostEnumerator.enumerateProgram(prog);
 		}
 		catch (IOException e) {
 			e.printStackTrace();

src/test/scripts/functions/federated/privacy/FederatedPlanCostEnumeratorTest4.dml

@@ -0,0 +1,28 @@
+#-------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+# 
+#   http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+#-------------------------------------------------------------
+
+a = matrix(7,10,10);
+if (sum(a) > 0.5)
+    b = a * 2;
+else
+    b = a * 3;
+c = sqrt(b);
+print(sum(c));

src/test/scripts/functions/federated/privacy/FederatedPlanCostEnumeratorTest5.dml

@@ -0,0 +1,26 @@
+#-------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+# 
+#   http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+#-------------------------------------------------------------
+
+for( i in 1:10 )
+{
+   b = i + 1;
+   print(b);
+}