Update and method cleanup

apache · Jul 17, 2022 · 25d8e83 · 25d8e83
1 parent 773752e
commit 25d8e83
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 16 deletions.
diff --git a/src/main/java/org/apache/sysds/hops/fedplanner/FederatedPlannerCostbased.java b/src/main/java/org/apache/sysds/hops/fedplanner/FederatedPlannerCostbased.java
@@ -197,23 +197,15 @@ private ArrayList<StatementBlock> rewriteDefaultStatementBlock(DMLProgram prog,
 					FunctionStatementBlock sbFuncBlock = prog.getFunctionDictionary(funcNamespace).getFunction(funcName);
 					FunctionStatement funcStatement = (FunctionStatement) sbFuncBlock.getStatement(0);
 
-					paramMap = createFunctionFedVarTable(paramMap, (FunctionOp) sbHop);
+					paramMap = FederatedPlannerUtils.createFunctionFedVarTable(paramMap, (FunctionOp) sbHop);
 					rewriteStatementBlock(prog, sbFuncBlock, paramMap);
-					FederatedPlannerUtils.mapFunctionOutputs((FunctionOp) sbHop, funcStatement);
+					FederatedPlannerUtils.mapFunctionOutputs((FunctionOp) sbHop, funcStatement, transientWrites);
 				}
 			}
 		}
 		return new ArrayList<>(Collections.singletonList(sb));
 	}
 
-	private Map<String, Hop> createFunctionFedVarTable(Map<String, Hop> paramMap, FunctionOp sbHop) {
-		Map<String, Hop> funcParamMap = FederatedPlannerUtils.getParamMap(sbHop);
-		if ( paramMap != null && funcParamMap != null)
-			funcParamMap.putAll(paramMap);
-		paramMap = funcParamMap;
-		return paramMap;
-	}
-
 	/**
 	 * Set final fedouts of all hops starting from terminal hops.
 	 */

diff --git a/src/main/java/org/apache/sysds/hops/fedplanner/FederatedPlannerFedAll.java b/src/main/java/org/apache/sysds/hops/fedplanner/FederatedPlannerFedAll.java
@@ -19,7 +19,6 @@
 
 package org.apache.sysds.hops.fedplanner;
 
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -46,7 +45,6 @@
 import org.apache.sysds.runtime.controlprogram.caching.CacheableData;
 import org.apache.sysds.runtime.instructions.cp.Data;
 import org.apache.sysds.runtime.instructions.fed.FEDInstruction.FederatedOutput;
-import org.jetbrains.annotations.NotNull;
 
 /**
  * Baseline federated planner that compiles all hops
@@ -166,7 +164,8 @@ else if( allowsFederated(hop, memo) ) {
 			memo.put(hop.getHopID(), null);
 	}
 
-	@NotNull
+	// TODO: Reduce code duplication. See `createFunctionFedVarTable` and `mapFunctionOutputs` in
+	//  `FederationPlannerUtils.java`
 	static private Map<String, FType> createFunctionFedVarTable(FunctionOp hop, Map<Long, FType> memo) {
 		Map<String, Hop> funcParamMap = FederatedPlannerUtils.getParamMap(hop);
 		Map<String, FType> funcFedVars = new HashMap<>();
@@ -176,7 +175,6 @@ static private Map<String, FType> createFunctionFedVarTable(FunctionOp hop, Map<
 		return funcFedVars;
 	}
 
-	// TODO: Reduce code duplication. The general structure of the federated planners is similar, but memo tables are different.
 	private void mapFunctionOutputs(FunctionOp sbHop, FunctionStatement funcStatement,
 		Map<String, FType> funcFedVars, Map<String, FType> callFedVars) {
 		for(int i = 0; i < sbHop.getOutputVariableNames().length; ++i) {

diff --git a/src/main/java/org/apache/sysds/hops/fedplanner/FederatedPlannerUtils.java b/src/main/java/org/apache/sysds/hops/fedplanner/FederatedPlannerUtils.java
@@ -92,10 +92,19 @@ public static Map<String,Hop> getParamMap(FunctionOp funcOp){
 	 * @param funcStatement The <code>FunctionStatement</code> of the called function
 	 * @param transientWrites map of transient writes
 	 */
-	public static void mapFunctionOutputs(FunctionOp sbHop, FunctionStatement funcStatement, Map<String,Hop> transientWrites) {
-		for (int i = 0; i < sbHop.getOutputVariableNames().length; ++i) {
+	public static void mapFunctionOutputs(FunctionOp sbHop, FunctionStatement funcStatement,
+		Map<String, Hop> transientWrites) {
+		for(int i = 0; i < sbHop.getOutputVariableNames().length; ++i) {
 			Hop outputWrite = transientWrites.get(funcStatement.getOutputParams().get(i).getName());
 			transientWrites.put(sbHop.getOutputVariableNames()[i], outputWrite);
 		}
 	}
+
+	public static Map<String, Hop> createFunctionFedVarTable(Map<String, Hop> paramMap, FunctionOp sbHop) {
+		Map<String, Hop> funcParamMap = FederatedPlannerUtils.getParamMap(sbHop);
+		if(paramMap != null)
+			funcParamMap.putAll(paramMap);
+		paramMap = funcParamMap;
+		return paramMap;
+	}
 }