Some more improvements

Author: janniklinde

src/main/java/org/apache/sysds/hops/rewriter/RewriterRuleCollection.java

@@ -471,6 +471,14 @@ public static void substituteFusedOps(final List<RewriterRule> rules, final Rule
 				.build()
 		);
 
+		rules.add(new RewriterRuleBuilder(ctx, "*2(A) => +(A,A)")
+				.setUnidirectional(true)
+				.parseGlobalVars("MATRIX:A")
+				.withParsedStatement("*2(A)")
+				.toParsedStatement("+(A,A)")
+				.build()
+		);
+
 		// TODO
 		/*rules.add(new RewriterRuleBuilder(ctx, "replace(A, a, b) => A")
 				.setUnidirectional(true)

src/main/java/org/apache/sysds/hops/rewriter/RewriterUtils.java

@@ -1461,6 +1461,8 @@ private static RewriterStatement tryPullOutSum(RewriterStatement sum, final Rule
 				return RewriterStatement.multiArgInstr(ctx, "*", toRemove.toArray(RewriterStatement[]::new));
 			}
 		} else if (sumBody.trueInstruction().equals("+")) {
+			// TODO: What about sum(+(A, *(a, B)))? We could pull out a
+
 			// We have to assume here, that this instruction is not referenced anywhere else in the graph
 			List<RewriterStatement> argList = sumBody.getChild(0).getOperands();
 			List<RewriterStatement> toRemove = new ArrayList<>(argList.size());
@@ -1489,9 +1491,10 @@ private static RewriterStatement tryPullOutSum(RewriterStatement sum, final Rule
 				}
 
 				mul.add(outerSum);
-				mul.add(sum);
+				RewriterStatement mulStmt = RewriterStatement.multiArgInstr(ctx, "*", mul.toArray(RewriterStatement[]::new));
+				//mul.add(sum);
 
-				return RewriterStatement.multiArgInstr(ctx, "*", mul.toArray(RewriterStatement[]::new));
+				return RewriterStatement.multiArgInstr(ctx, "+", mulStmt, sum);
 			}
 		}
 

src/test/java/org/apache/sysds/test/component/codegen/rewrite/RewriterStreamTests.java

@@ -193,14 +193,19 @@ public void testSumEquality6() {
 	@Test
 	public void testSumEquality() {
 		RewriterStatement stmt = RewriterUtils.parse("sum(+(B, sum(*(a, A))))", ctx, "MATRIX:A,B", "FLOAT:a");
-		RewriterStatement stmt2 = RewriterUtils.parse("*(a, sum(+(B, sum(A))))", ctx, "MATRIX:A,B", "FLOAT:a");
+		RewriterStatement stmt2 = RewriterUtils.parse("+(*(a, length(A)), sum(+(B, sum(A))))", ctx, "MATRIX:A,B", "FLOAT:a");
+		RewriterStatement stmt3 = RewriterUtils.parse("sum(+(B, *(a, sum(A))))", ctx, "MATRIX:A,B", "FLOAT:a");
 		stmt = canonicalConverter.apply(stmt);
 		stmt2 = canonicalConverter.apply(stmt2);
+		stmt3 = canonicalConverter.apply(stmt3);
 
 		System.out.println("==========");
 		System.out.println(stmt.toParsableString(ctx, true));
 		System.out.println("==========");
+		System.out.println(stmt3.toParsableString(ctx, true));
+		System.out.println("==========");
 		System.out.println(stmt2.toParsableString(ctx, true));
+		assert stmt.match(RewriterStatement.MatcherContext.exactMatch(ctx, stmt3, stmt));
 		assert stmt.match(RewriterStatement.MatcherContext.exactMatch(ctx, stmt2, stmt));
 	}
 
@@ -1245,4 +1250,42 @@ public void testFused6() {
 
 		assert !stmt1.match(RewriterStatement.MatcherContext.exactMatch(ctx, stmt2, stmt1));
 	}
+
+	@Test
+	public void testSum() {
+		RewriterStatement stmt1 = RewriterUtils.parse("sum(+(a,A))", ctx, "MATRIX:A,B", "FLOAT:a");
+		RewriterStatement stmt2 = RewriterUtils.parse("+(*(a, length(A)), sum(A))", ctx, "MATRIX:A,B", "FLOAT:a", "LITERAL_FLOAT:0.0");
+
+		System.out.println("Cost1: " + RewriterCostEstimator.estimateCost(stmt1, ctx));
+		System.out.println("Cost2: " + RewriterCostEstimator.estimateCost(stmt2, ctx));
+
+		stmt1 = canonicalConverter.apply(stmt1);
+		stmt2 = canonicalConverter.apply(stmt2);
+
+		System.out.println("==========");
+		System.out.println(stmt1.toParsableString(ctx, true));
+		System.out.println("==========");
+		System.out.println(stmt2.toParsableString(ctx, true));
+
+		assert stmt1.match(RewriterStatement.MatcherContext.exactMatch(ctx, stmt2, stmt1));
+	}
+
+	@Test
+	public void testSumInequality() {
+		RewriterStatement stmt1 = RewriterUtils.parse("sum(+(a,*(B,c)))", ctx, "MATRIX:B", "FLOAT:a,c");
+		RewriterStatement stmt2 = RewriterUtils.parse("*(a, sum(+(B,c)))", ctx, "MATRIX:B", "FLOAT:a,c", "LITERAL_FLOAT:0.0");
+
+		System.out.println("Cost1: " + RewriterCostEstimator.estimateCost(stmt1, ctx));
+		System.out.println("Cost2: " + RewriterCostEstimator.estimateCost(stmt2, ctx));
+
+		stmt1 = canonicalConverter.apply(stmt1);
+		stmt2 = canonicalConverter.apply(stmt2);
+
+		System.out.println("==========");
+		System.out.println(stmt1.toParsableString(ctx, true));
+		System.out.println("==========");
+		System.out.println(stmt2.toParsableString(ctx, true));
+
+		assert !stmt1.match(RewriterStatement.MatcherContext.exactMatch(ctx, stmt2, stmt1));
+	}
 }

src/test/java/org/apache/sysds/test/component/codegen/rewrite/functions/DMLCodeGenTest.java

@@ -148,4 +148,21 @@ public void testFused2() {
 
 		assert RewriterRuleCreator.validateRuleApplicability(rule, ctx);
 	}
+
+	@Test
+	public void testFused3() {
+		// TODO: This rule has been ignored, but why?
+		String ruleStr = "MATRIX:A,B\nLITERAL_FLOAT:0.0,1.0\n" +
+				"+(-(A,B),A)\n" +
+				"=>\n" +
+				"-(*2(A), B)";
+
+		RewriterRule rule = RewriterUtils.parseRule(ruleStr, ctx);
+
+		System.out.println(DMLCodeGenerator.generateRuleValidationDML(rule, "test", ctx));
+
+		assert RewriterRuleCreator.validateRuleCorrectness(rule, ctx);
+
+		assert RewriterRuleCreator.validateRuleApplicability(rule, ctx);
+	}
 }