[SYSTEMDS-3941] Add new rewrites

src/main/java/org/apache/sysds/hops/rewrite/RewriteAlgebraicSimplificationDynamic.java

@@ -184,6 +184,7 @@ private void rule_AlgebraicSimplification(Hop hop, boolean descendFirst)
 			hi = simplifySumDiagToTrace(hi);                  //e.g., sum(diag(X)) -> trace(X); if col vector
 			hi = simplifyLowerTriExtraction(hop, hi, i);      //e.g., X * cumsum(diag(matrix(1,nrow(X),1))) -> lower.tri
 			hi = simplifyConstantCumsum(hop, hi, i);          //e.g., cumsum(matrix(1/n,n,1)) -> seq(1/n, 1, 1/n)
+			hi = simplifySumConstantMatrix(hop, hi, i);       //e.g., sum(matrix(a,rows=b,cols=c)) -> a*b*c
 			hi = pushdownBinaryOperationOnDiag(hop, hi, i);   //e.g., diag(X)*7 -> diag(X*7); if col vector
 			hi = pushdownSumOnAdditiveBinary(hop, hi, i);     //e.g., sum(A+B) -> sum(A)+sum(B); if dims(A)==dims(B)
 			if(OptimizerUtils.ALLOW_OPERATOR_FUSION) {
@@ -1273,6 +1274,31 @@ private static Hop simplifyConstantCumsum(Hop parent, Hop hi, int pos) {
 		}
 		return hi;
 	}
+
+	private static Hop simplifySumConstantMatrix(Hop parent, Hop hi, int pos) {
+		//pattern: sum(matrix(a, rows=b, cols=c)) -> a*b*c
+		if( HopRewriteUtils.isAggUnaryOp(hi, AggOp.SUM, Direction.RowCol)
+			&& HopRewriteUtils.isDataGenOpWithConstantValue(hi.getInput(0))
+			&& hi.getInput(0).dimsKnown()
+			&& hi.getInput(0).getDim1() >= 1
+			&& hi.getInput(0).getDim2() >= 1
+			&& hi.getInput(0).getParent().size() == 1 )
+		{
+			DataGenOp datagen = (DataGenOp) hi.getInput(0);
+			Hop constVal = datagen.getConstantValue();
+			Hop rows = new LiteralOp(datagen.getDim1());
+			Hop cols = new LiteralOp(datagen.getDim2());
+
+			Hop hnew = HopRewriteUtils.createBinary(
+				HopRewriteUtils.createBinary(constVal, rows, OpOp2.MULT), cols, OpOp2.MULT);
+			HopRewriteUtils.replaceChildReference(parent, hi, hnew, pos);
+			HopRewriteUtils.cleanupUnreferenced(hi, datagen);
+
+			hi = hnew;
+			LOG.debug("Applied simplifySumConstantMatrix (line "+hi.getBeginLine()+").");
+		}
+		return hi;
+	}
 
 	private static Hop pushdownBinaryOperationOnDiag(Hop parent, Hop hi, int pos) 
 	{

src/main/java/org/apache/sysds/hops/rewrite/RewriteAlgebraicSimplificationStatic.java

@@ -170,6 +170,8 @@ private void rule_AlgebraicSimplification(Hop hop, boolean descendFirst)
 			hi = pushdownDetMultOperation(hop, hi, i);           //e.g., det(X%*%Y) -> det(X)*det(Y)
 			hi = pushdownDetScalarMatrixMultOperation(hop, hi, i);  //e.g., det(lambda*X) -> lambda^nrow(X)*det(X)
 			hi = pushdownSumBinaryMult(hop, hi, i);              //e.g., sum(lambda*X) -> lambda*sum(X)
+			hi = pushdownRowSumBinaryMult(hop, hi, i);           //e.g., rowSums(lambda*X) -> lambda*rowSums(X)
+			hi = pushdownColSumBinaryMult(hop, hi, i);           //e.g., colSums(lambda*X) -> lambda*colSums(X)
 			hi = pullupAbs(hop, hi, i);                          //e.g., abs(X)*abs(Y) --> abs(X*Y)
 			hi = simplifyUnaryPPredOperation(hop, hi, i);        //e.g., abs(ppred()) -> ppred(), others: round, ceil, floor
 			hi = simplifyTransposedAppend(hop, hi, i);           //e.g., t(cbind(t(A),t(B))) -> rbind(A,B);
@@ -1447,6 +1449,58 @@ private static Hop pushdownSumBinaryMult(Hop parent, Hop hi, int pos ) {
 		return hi;
 	}
 
+	private static Hop pushdownRowSumBinaryMult(Hop parent, Hop hi, int pos ) {
+		//pattern:  rowSums(lamda*X) -> lamda*rowSums(X)
+		if( hi instanceof AggUnaryOp && ((AggUnaryOp)hi).getDirection()==Direction.Row
+				&& ((AggUnaryOp)hi).getOp()==AggOp.SUM // only one parent which is the rowSums
+				&& HopRewriteUtils.isBinary(hi.getInput(0), OpOp2.MULT, 1)
+				&& ((hi.getInput(0).getInput(0).getDataType()==DataType.SCALAR && hi.getInput(0).getInput(1).getDataType()==DataType.MATRIX)
+				||(hi.getInput(0).getInput(0).getDataType()==DataType.MATRIX && hi.getInput(0).getInput(1).getDataType()==DataType.SCALAR)))
+		{
+			Hop operand1 = hi.getInput(0).getInput(0);
+			Hop operand2 = hi.getInput(0).getInput(1);
+
+			//check which operand is the Scalar and which is the matrix
+			Hop lamda = (operand1.getDataType()==DataType.SCALAR) ? operand1 : operand2;
+			Hop matrix = (operand1.getDataType()==DataType.MATRIX) ? operand1 : operand2;
+
+			AggUnaryOp aggOp=HopRewriteUtils.createAggUnaryOp(matrix, AggOp.SUM, Direction.Row);
+			Hop bop = HopRewriteUtils.createBinary(lamda, aggOp, OpOp2.MULT);
+
+			HopRewriteUtils.replaceChildReference(parent, hi, bop, pos);
+
+			LOG.debug("Applied pushdownRowSumBinaryMult (line "+hi.getBeginLine()+").");
+			return bop;
+		}
+		return hi;
+	}
+
+	private static Hop pushdownColSumBinaryMult(Hop parent, Hop hi, int pos ) {
+		//pattern:  colSums(lamda*X) -> lamda*colSums(X)
+		if( hi instanceof AggUnaryOp && ((AggUnaryOp)hi).getDirection()==Direction.Col
+				&& ((AggUnaryOp)hi).getOp()==AggOp.SUM // only one parent which is the colSums
+				&& HopRewriteUtils.isBinary(hi.getInput(0), OpOp2.MULT, 1)
+				&& ((hi.getInput(0).getInput(0).getDataType()==DataType.SCALAR && hi.getInput(0).getInput(1).getDataType()==DataType.MATRIX)
+				||(hi.getInput(0).getInput(0).getDataType()==DataType.MATRIX && hi.getInput(0).getInput(1).getDataType()==DataType.SCALAR)))
+		{
+			Hop operand1 = hi.getInput(0).getInput(0);
+			Hop operand2 = hi.getInput(0).getInput(1);
+
+			//check which operand is the Scalar and which is the matrix
+			Hop lamda = (operand1.getDataType()==DataType.SCALAR) ? operand1 : operand2;
+			Hop matrix = (operand1.getDataType()==DataType.MATRIX) ? operand1 : operand2;
+
+			AggUnaryOp aggOp=HopRewriteUtils.createAggUnaryOp(matrix, AggOp.SUM, Direction.Col);
+			Hop bop = HopRewriteUtils.createBinary(lamda, aggOp, OpOp2.MULT);
+
+			HopRewriteUtils.replaceChildReference(parent, hi, bop, pos);
+
+			LOG.debug("Applied pushdownColSumBinaryMult (line "+hi.getBeginLine()+").");
+			return bop;
+		}
+		return hi;
+	}
+
 	private static Hop pullupAbs(Hop parent, Hop hi, int pos ) {
 		if( HopRewriteUtils.isBinary(hi, OpOp2.MULT)
 				&& HopRewriteUtils.isUnary(hi.getInput(0), OpOp1.ABS)

src/test/java/org/apache/sysds/test/functions/rewrite/RewriteFusedRandTest.java

@@ -121,7 +121,7 @@ private void testRewriteFusedRand( String testname, String pdf, boolean rewrites
 			//compare matrices 
 			Double ret = readDMLMatrixFromOutputDir("R").get(new CellIndex(1,1));
 			if( testname.equals(TEST_NAME1) )
-				Assert.assertEquals("Wrong result", Double.valueOf(rows), ret);
+				Assert.assertEquals("Wrong result", Double.valueOf(rows*cols), ret);
 			else if( testname.equals(TEST_NAME2) )
 				Assert.assertEquals("Wrong result", Double.valueOf(Math.pow(rows*cols, 2)), ret);
 			else if( testname.equals(TEST_NAME3) )

src/test/java/org/apache/sysds/test/functions/rewrite/RewritePushdownColSumBinaryMultTest.java

@@ -0,0 +1,100 @@
+/*
+ * 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.
+ */
+
+package org.apache.sysds.test.functions.rewrite;
+
+import java.util.HashMap;
+
+import org.junit.Test;
+import org.apache.sysds.hops.OptimizerUtils;
+import org.apache.sysds.runtime.matrix.data.MatrixValue.CellIndex;
+import org.apache.sysds.test.AutomatedTestBase;
+import org.apache.sysds.test.TestConfiguration;
+import org.apache.sysds.test.TestUtils;
+import org.apache.sysds.utils.Statistics;
+import org.junit.Assert;
+
+public class RewritePushdownColSumBinaryMultTest extends AutomatedTestBase
+{
+	private static final String TEST_NAME1 = "RewritePushdownColSumBinaryMult";
+	private static final String TEST_NAME2 = "RewritePushdownColSumBinaryMult2";
+
+	private static final String TEST_DIR = "functions/rewrite/";
+	private static final String TEST_CLASS_DIR = TEST_DIR + RewritePushdownColSumBinaryMultTest.class.getSimpleName() + "/";
+
+	@Override
+	public void setUp() {
+		TestUtils.clearAssertionInformation();
+		addTestConfiguration(TEST_NAME1, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME1, new String[] { "R" }));
+		addTestConfiguration(TEST_NAME2, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME2, new String[] { "R" }));
+	}
+
+	@Test
+	public void testPushdownColSumBinaryMultNoRewrite() {
+		testRewritePushdownColSumBinaryMult(TEST_NAME1, false);
+	}
+
+	@Test
+	public void testPushdownColSumBinaryMultRewrite() {
+		testRewritePushdownColSumBinaryMult(TEST_NAME1, true);
+	}
+
+	@Test
+	public void testPushdownColSumBinaryMultNoRewrite2() {
+		testRewritePushdownColSumBinaryMult(TEST_NAME2, false);
+	}
+
+	@Test
+	public void testPushdownColSumBinaryMultRewrite2() {
+		testRewritePushdownColSumBinaryMult(TEST_NAME2, true);
+	}
+
+	private void testRewritePushdownColSumBinaryMult(String testname, boolean rewrites) {
+		boolean oldFlag = OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION;
+
+		try {
+			TestConfiguration config = getTestConfiguration(testname);
+			loadTestConfiguration(config);
+
+			String HOME = SCRIPT_DIR + TEST_DIR;
+			fullDMLScriptName = HOME + testname + ".dml";
+			programArgs = new String[] { "-stats", "-args", output("R") };
+
+			fullRScriptName = HOME + testname + ".R";
+			rCmd = getRCmd(inputDir(), expectedDir());
+
+			OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION = rewrites;
+
+			runTest(true, false, null, -1);
+			runRScript(true);
+
+			HashMap<CellIndex, Double> dmlfile = readDMLMatrixFromOutputDir("R");
+			HashMap<CellIndex, Double> rfile = readRMatrixFromExpectedDir("R");
+			TestUtils.compareMatrices(dmlfile, rfile, 1e-10, "DML", "R");
+
+			if(rewrites)
+				Assert.assertEquals(1, Statistics.getCPHeavyHitterCount("n*"));
+			else
+				Assert.assertEquals(2, Statistics.getCPHeavyHitterCount("*"));
+		}
+		finally {
+			OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION = oldFlag;
+		}
+	}
+}

src/test/java/org/apache/sysds/test/functions/rewrite/RewritePushdownRowSumBinaryMultTest.java

@@ -0,0 +1,100 @@
+/*
+ * 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.
+ */
+
+package org.apache.sysds.test.functions.rewrite;
+
+import java.util.HashMap;
+
+import org.junit.Test;
+import org.apache.sysds.hops.OptimizerUtils;
+import org.apache.sysds.runtime.matrix.data.MatrixValue.CellIndex;
+import org.apache.sysds.test.AutomatedTestBase;
+import org.apache.sysds.test.TestConfiguration;
+import org.apache.sysds.test.TestUtils;
+import org.apache.sysds.utils.Statistics;
+import org.junit.Assert;
+
+public class RewritePushdownRowSumBinaryMultTest extends AutomatedTestBase
+{
+	private static final String TEST_NAME1 = "RewritePushdownRowSumBinaryMult";
+	private static final String TEST_NAME2 = "RewritePushdownRowSumBinaryMult2";
+
+	private static final String TEST_DIR = "functions/rewrite/";
+	private static final String TEST_CLASS_DIR = TEST_DIR + RewritePushdownRowSumBinaryMultTest.class.getSimpleName() + "/";
+
+	@Override
+	public void setUp() {
+		TestUtils.clearAssertionInformation();
+		addTestConfiguration(TEST_NAME1, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME1, new String[] { "R" }));
+		addTestConfiguration(TEST_NAME2, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME2, new String[] { "R" }));
+	}
+
+	@Test
+	public void testPushdownRowSumBinaryMultNoRewrite() {
+		testRewritePushdownRowSumBinaryMult(TEST_NAME1, false);
+	}
+
+	@Test
+	public void testPushdownRowSumBinaryMultRewrite() {
+		testRewritePushdownRowSumBinaryMult(TEST_NAME1, true);
+	}
+
+	@Test
+	public void testPushdownRowSumBinaryMultNoRewrite2() {
+		testRewritePushdownRowSumBinaryMult(TEST_NAME2, false);
+	}
+
+	@Test
+	public void testPushdownRowSumBinaryMultRewrite2() {
+		testRewritePushdownRowSumBinaryMult(TEST_NAME2, true);
+	}
+
+	private void testRewritePushdownRowSumBinaryMult(String testname, boolean rewrites) {
+		boolean oldFlag = OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION;
+
+		try {
+			TestConfiguration config = getTestConfiguration(testname);
+			loadTestConfiguration(config);
+
+			String HOME = SCRIPT_DIR + TEST_DIR;
+			fullDMLScriptName = HOME + testname + ".dml";
+			programArgs = new String[] { "-stats", "-args", output("R") };
+
+			fullRScriptName = HOME + testname + ".R";
+			rCmd = getRCmd(inputDir(), expectedDir());
+
+			OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION = rewrites;
+
+			runTest(true, false, null, -1);
+			runRScript(true);
+
+			HashMap<CellIndex, Double> dmlfile = readDMLMatrixFromOutputDir("R");
+			HashMap<CellIndex, Double> rfile = readRMatrixFromExpectedDir("R");
+			TestUtils.compareMatrices(dmlfile, rfile, 1e-10, "DML", "R");
+
+			if(rewrites)
+				Assert.assertEquals(1, Statistics.getCPHeavyHitterCount("n*"));
+			else
+				Assert.assertEquals(2, Statistics.getCPHeavyHitterCount("*"));
+		}
+		finally {
+			OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION = oldFlag;
+		}
+	}
+}