Some major performance improvements

Jaybit0 · Jan 29, 2025 · 31f8029 · 31f8029
1 parent 0e15664
commit 31f8029
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 9 deletions.
diff --git a/src/main/java/org/apache/sysds/hops/rewriter/RewriterDataType.java b/src/main/java/org/apache/sysds/hops/rewriter/RewriterDataType.java
@@ -30,6 +30,9 @@ public String getResultingDataType(final RuleContext ctx) {
 		return type;
 	}
 
+	@Override
+	public void refreshReturnType(final RuleContext ctx) {}
+
 	@Override
 	public boolean isLiteral() {
 		return literal != null && !(literal instanceof List<?>);

diff --git a/src/main/java/org/apache/sysds/hops/rewriter/RewriterInstruction.java b/src/main/java/org/apache/sysds/hops/rewriter/RewriterInstruction.java
@@ -18,26 +18,35 @@
 public class RewriterInstruction extends RewriterStatement {
 
 	private String id;
+	private String returnType;
 	private String instr;
 	//private RewriterDataType result = new RewriterDataType();
 	private ArrayList<RewriterStatement> operands = new ArrayList<>();
 	private Function<List<RewriterStatement>, Long> costFunction = null;
 	private boolean consolidated = false;
 	private int hashCode;
 
-	//private DualHashBidiMap<RewriterStatement, RewriterStatement> links = null;
-
 	@Override
 	public String getId() {
 		return id;
 	}
 
 	@Override
 	public String getResultingDataType(final RuleContext ctx) {
-		if (isArgumentList()) {
-			return getOperands().stream().map(op -> op.getResultingDataType(ctx)).reduce(RewriterUtils::defaultTypeHierarchy).get() + "...";
-		}
-		return ctx.instrTypes.get(trueTypedInstruction(ctx));//getResult(ctx).getResultingDataType(ctx);
+		if (returnType != null)
+			return returnType;
+
+		if (isArgumentList())
+			returnType = getOperands().stream().map(op -> op.getResultingDataType(ctx)).reduce(RewriterUtils::defaultTypeHierarchy).get() + "...";
+		else
+			returnType = ctx.instrTypes.get(trueTypedInstruction(ctx));//getResult(ctx).getResultingDataType(ctx);
+
+		return returnType;
+	}
+
+	@Override
+	public void refreshReturnType(final RuleContext ctx) {
+		returnType = null;
 	}
 
 	@Override
@@ -472,9 +481,10 @@ public boolean hasProperty(String property, final RuleContext ctx) {
 	}
 
 	public String trueInstruction() {
-		Object trueInstrObj = getMeta("trueInstr");
+		// Legacy code
+		/*Object trueInstrObj = getMeta("trueInstr");
 		if (trueInstrObj != null && trueInstrObj instanceof String)
-			return (String)trueInstrObj;
+			return (String)trueInstrObj;*/
 		return instr;
 	}
 

diff --git a/src/main/java/org/apache/sysds/hops/rewriter/RewriterRuleCollection.java b/src/main/java/org/apache/sysds/hops/rewriter/RewriterRuleCollection.java
@@ -1179,7 +1179,7 @@ public static void canonicalExpandAfterFlattening(final List<RewriterRule> rules
 				.parseGlobalVars("INT...:indices")
 				.parseGlobalVars("FLOAT...:ops")
 				.withParsedStatement("sum($1:_idxExpr(indices, +(ops)))", hooks)
-				.toParsedStatement("+($3:argList(sum($2:_idxExpr(indices, +(ops)))))", hooks) // The inner +(ops) is temporary and will be removed
+				.toParsedStatement("$4:+($3:argList(sum($2:_idxExpr(indices, +(ops)))))", hooks) // The inner +(ops) is temporary and will be removed
 				.link(hooks.get(1).getId(), hooks.get(2).getId(), RewriterStatement::transferMeta)
 				.apply(hooks.get(3).getId(), newArgList -> {
 					RewriterStatement oldArgList = newArgList.getChild(0, 0, 1, 0);
@@ -1196,6 +1196,11 @@ public static void canonicalExpandAfterFlattening(final List<RewriterRule> rules
 						RewriterUtils.copyIndexList(newIdxExpr);
 						newArgList.getOperands().add(newSum);
 					}
+
+					newArgList.refreshReturnType(ctx);
+				}, true)
+				.apply(hooks.get(4).getId(), stmt -> {
+					stmt.refreshReturnType(ctx);
 				}, true)
 				.build()
 		);

diff --git a/src/main/java/org/apache/sysds/hops/rewriter/RewriterStatement.java b/src/main/java/org/apache/sysds/hops/rewriter/RewriterStatement.java
@@ -296,6 +296,7 @@ public void setLiteral(Object literal) {
 	public abstract RewriterStatement nestedCopyOrInject(Map<RewriterStatement, RewriterStatement> copiedObjects, TriFunction<RewriterStatement, RewriterStatement, Integer, RewriterStatement> injector, RewriterStatement parent, int pIdx);
 	// Returns the new maxRefId
 	abstract int toParsableString(StringBuilder builder, Map<RewriterRule.IdentityRewriterStatement, Integer> refs, int maxRefId, Map<String, Set<String>> vars, final RuleContext ctx);
+	abstract void refreshReturnType(final RuleContext ctx);
 
 	public String toParsableString(final RuleContext ctx, boolean includeDefinitions) {
 		StringBuilder sb = new StringBuilder();

diff --git a/src/main/java/org/apache/sysds/hops/rewriter/RewriterUtils.java b/src/main/java/org/apache/sysds/hops/rewriter/RewriterUtils.java
@@ -279,6 +279,7 @@ public static void mergeArgLists(RewriterStatement stmt, final RuleContext ctx)
 		stmt.forEachPreOrder(el -> {
 			tryFlattenNestedArgList(ctx, el, el, -1);
 			tryFlattenNestedOperatorPatterns(ctx, el);
+			el.refreshReturnType(ctx);
 			return true;
 		});
 

diff --git a/src/test/java/org/apache/sysds/test/component/codegen/rewrite/RewriterClusteringTest.java b/src/test/java/org/apache/sysds/test/component/codegen/rewrite/RewriterClusteringTest.java
@@ -1,5 +1,6 @@
 package org.apache.sysds.test.component.codegen.rewrite;
 
+import org.apache.commons.lang3.mutable.MutableInt;
 import org.apache.commons.lang3.mutable.MutableLong;
 import org.apache.sysds.hops.rewriter.RewriterDatabase;
 import org.apache.sysds.hops.rewriter.RewriterHeuristic;
@@ -61,9 +62,16 @@ public void testExpressionClustering() {
 		RewriterDatabase canonicalExprDB = new RewriterDatabase();
 		List<RewriterStatement> foundEquivalences = new ArrayList<>();
 
+		int size = db.size();
+		MutableInt ctr = new MutableInt(0);
+
 		db.forEach(expr -> {
+			if (ctr.incrementAndGet() % 10 == 0)
+				System.out.println("Done: " + ctr.intValue() + " / " + size);
 			// First, build all possible subtrees
 			List<RewriterStatement> subExprs = RewriterUtils.generateSubtrees(expr, ctx);
+			if (subExprs.size() > 100)
+				System.out.println("Critical number of subtrees: " + subExprs.size());
 			//List<RewriterStatement> subExprs = List.of(expr);
 			long evaluationCtr = 0;