Some progress

Author: janniklinde

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

@@ -765,7 +765,7 @@ public static void expandStreamingExpressions(final List<RewriterRule> rules, fi
 				.parseGlobalVars("LITERAL_INT:1")
 				.parseGlobalVars("LITERAL_FLOAT:0.0")
 				.withParsedStatement("diag(A)", hooks)
-				.toParsedStatement("$4:_m($1:_idx(1, nrow(A)), $2:_idx(1, ncol(A)), [](A, $1, $1))", hooks)
+				.toParsedStatement("$4:_m($1:_idx(1, nrow(A)), $2:_idx(1, ncol(A)), ifelse(==($1,$2), [](A, $1, $2), 0.0))", hooks)
 				.apply(hooks.get(1).getId(), stmt -> stmt.unsafePutMeta("idxId", UUID.randomUUID()), true) // Assumes it will never collide
 				.apply(hooks.get(2).getId(), stmt -> stmt.unsafePutMeta("idxId", UUID.randomUUID()), true) // Assumes it will never collide
 				.apply(hooks.get(4).getId(), (stmt, match) -> {
@@ -775,8 +775,8 @@ public static void expandStreamingExpressions(final List<RewriterRule> rules, fi
 
 					RewriterStatement aRef = stmt.getChild(0, 1, 0);
 
-					System.out.println("GETTING: ");
-					System.out.println(match.getNewExprRoot().getAssertions(ctx).getAssertionStatement(aRef.getNCol(), null));
+					//System.out.println("GETTING: ");
+					//System.out.println(match.getNewExprRoot().getAssertions(ctx).getAssertionStatement(aRef.getNCol(), null));
 					match.getNewExprRoot().getAssertions(ctx).addEqualityAssertion(aRef.getNCol(), aRef.getNRow(), match.getNewExprRoot());
 				}, true) // Assumes it will never collide
 				.build()
@@ -1163,6 +1163,18 @@ public static void expandArbitraryMatrices(final List<RewriterRule> rules, final
 	public static void pushdownStreamSelections(final List<RewriterRule> rules, final RuleContext ctx) {
 		HashMap<Integer, RewriterStatement> hooks = new HashMap<>();
 
+		// ifelse merging
+		// TODO: Permutations e.g. ==(l2, l1) etc.
+		rules.add(new RewriterRuleBuilder(ctx)
+				.setUnidirectional(true)
+				.parseGlobalVars("FLOAT:a,b,c,d")
+				.parseGlobalVars("INT:l1,l2")
+				.withParsedStatement("$1:ElementWiseInstruction(ifelse(==(l1, l2), a, b), ifelse(==(l1, l2), c, d))", hooks)
+				.toParsedStatement("ifelse(==(l1, l2), $2:ElementWiseInstruction(a, c), $3:ElementWiseInstruction(b, d))", hooks)
+				.linkManyUnidirectional(hooks.get(1).getId(), List.of(hooks.get(2).getId(), hooks.get(3).getId()), RewriterStatement::transferMeta, true)
+				.build()
+		);
+
 		rules.add(new RewriterRuleBuilder(ctx)
 				.setUnidirectional(true)
 				.parseGlobalVars("INT:l")

src/main/java/org/apache/sysds/hops/rewriter/assertions/RewriterAssertions.java

@@ -466,11 +466,11 @@ public boolean addEqualityAssertion(RewriterStatement stmt1, RewriterStatement s
 
 				//System.out.println("MNew parts: " + partOfAssertion);
 
-				System.out.println("New assertion1: " + newAssertion);
+				//System.out.println("New assertion1: " + newAssertion);
 				return true;
 			}
 
-			System.out.println("Assertion already exists");
+			//System.out.println("Assertion already exists");
 			return false; // The assertion already exists
 		}
 
@@ -497,7 +497,7 @@ public boolean addEqualityAssertion(RewriterStatement stmt1, RewriterStatement s
 				return true;
 			}, false);
 
-			System.out.println("New assertion2: " + existingAssertion);
+			//System.out.println("New assertion2: " + existingAssertion);
 			return true;
 		}
 
@@ -521,7 +521,7 @@ public boolean addEqualityAssertion(RewriterStatement stmt1, RewriterStatement s
 		if (stmt1Assertions.stmt != null)
 			assertionMatcher.put(stmt1Assertions.stmt, stmt2Assertions); // Only temporary
 
-		System.out.println("New assertion3: " + stmt2Assertions);
+		//System.out.println("New assertion3: " + stmt2Assertions);
 		resolveCyclicAssertions(stmt2Assertions);
 		stmt2Assertions.deduplicate();
 

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

@@ -1432,5 +1432,6 @@ public void testDiag1() {
 		System.out.println(stmt2.toParsableString(ctx, true));
 
 		assert stmt1.match(RewriterStatement.MatcherContext.exactMatch(ctx, stmt2, stmt1));
+		assert RewriterCostEstimator.estimateCost(stmt1, ctx) > RewriterCostEstimator.estimateCost(stmt2, ctx);
 	}
 }