[SYSTEMDS-3986] Create class for Sparsity-based Optimization

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

@@ -149,6 +149,9 @@ public ProgramRewriter(boolean staticRewrites, boolean dynamicRewrites)
 			if(ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_REWRITES)) {
 				_dagRuleSet.add( new RewriteMatrixMultChainOptimizationSparse()         );
 			}
+			if( ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_RECOMPILE) ) {
+				_dagRuleSet.add(new SparsityOptimizer());
+			}
 			if( OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION ) {
 				_dagRuleSet.add( new RewriteAlgebraicSimplificationDynamic()      ); //dependencies: cse
 				_dagRuleSet.add( new RewriteAlgebraicSimplificationStatic()       ); //dependencies: cse

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

@@ -20,13 +20,10 @@
 package org.apache.sysds.hops.rewrite;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 
 import org.apache.sysds.common.Types.OpOpData;
-import org.apache.sysds.conf.ConfigurationManager;
-import org.apache.sysds.conf.DMLConfig;
 import org.apache.sysds.hops.AggBinaryOp;
 import org.apache.sysds.hops.DataOp;
 import org.apache.sysds.hops.Hop;
@@ -54,10 +51,6 @@ private static Hop getFirstTransientWriteParent(Hop hop) {
 	}
 
 	public List<StatementBlock> rRewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) {
-		// split only if sparsity rewrites are enabled (splitting is pointless w/o rewrite)
-		if(!ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_REWRITES))
-			return Arrays.asList(sb);
-
 		ArrayList<StatementBlock> ret = new ArrayList<>();
 
 		// collect all candidate hops in levels based on their data dependencies
@@ -158,23 +151,19 @@ private void collectNnzComputationCandidates(ArrayList<Hop> roots, ArrayList<Hop
 		}
 	}
 
-	public static boolean anyParentIsMatrixMultiply(Hop hop) {
+	private static boolean anyParentIsMatrixMultiply(Hop hop) {
 		return hop.getParent().stream().anyMatch(parentHop -> HopRewriteUtils.isMatrixMultiply(parentHop));
 	}
 
-	public static boolean allParentsAreMatrixMultiply(Hop hop) {
+	private static boolean allParentsAreMatrixMultiply(Hop hop) {
 		return hop.getParent().stream().allMatch(parentHop -> HopRewriteUtils.isMatrixMultiply(parentHop));
 	}
 
-	public static boolean anyChildIsMatrixMultiply(Hop hop) {
+	private static boolean anyChildIsMatrixMultiply(Hop hop) {
 		return hop.getInput().stream().anyMatch(childHop -> HopRewriteUtils.isMatrixMultiply(childHop));
 	}
 
-	public static boolean allChildsAreMatrixMultiply(Hop hop) {
-		return hop.getInput().stream().allMatch(childHop -> HopRewriteUtils.isMatrixMultiply(childHop));
-	}
-
-	public static boolean isMMChain(Hop hop) {
+	private static boolean isPartOfMMChain(Hop hop) {
 		// hop is part of mm chain with at least two matrix multplications
 		return HopRewriteUtils.isMatrixMultiply(hop)
 			&& (anyParentIsMatrixMultiply(hop) || anyChildIsMatrixMultiply(hop));
@@ -184,7 +173,7 @@ public static boolean anyParentIsMMChain(Hop hop) {
 		List<Hop> parentHops = hop.getParent();
 		if(parentHops == null)
 			return false;
-		return parentHops.stream().anyMatch(pHop -> isMMChain(pHop));
+		return parentHops.stream().anyMatch(pHop -> isPartOfMMChain(pHop));
 	}
 
 	private void rCollectNnzComputationCandidates(Hop hop, ArrayList<Hop> cand, Hop parentHop) {
@@ -205,7 +194,7 @@ && anyParentIsMMChain(hop)) { // split only mm chains (i.e., at least two matrix
 				cand.add(hop);
 			}
 			else if((!HopRewriteUtils.isMatrixMultiply(parentHop) || !allParentsAreMatrixMultiply(hop))
-				&& isMMChain(hop) // split only mm chains (i.e., at least two matrix multiplications)
+				&& isPartOfMMChain(hop) // split only mm chains (i.e., at least two matrix multiplications)
 				&& !((AggBinaryOp) hop).hasLeftPMInput()) {
 				// split after mm chain to force the mm chain into a separate statment block for isolated optimization
 				cand.add(hop);

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

@@ -0,0 +1,150 @@
+/*
+ * 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.hops.rewrite;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.apache.sysds.common.Types.OpOpData;
+import org.apache.sysds.hops.AggBinaryOp;
+import org.apache.sysds.hops.DataOp;
+import org.apache.sysds.hops.Hop;
+import org.apache.sysds.utils.Explain;
+
+public class SparsityOptimizer extends HopRewriteRule {
+	protected static final Log LOG = LogFactory.getLog(SparsityOptimizer.class.getName());
+
+	@Override
+	public ArrayList<Hop> rewriteHopDAGs(ArrayList<Hop> roots, ProgramRewriteStatus state) {
+		ArrayList<Hop> ret = new ArrayList<Hop>();
+		for(Hop rootHop : roots)
+			ret.add(rewriteHopDAG(rootHop, state));
+		return ret;
+	}
+
+	@Override
+	public Hop rewriteHopDAG(Hop root, ProgramRewriteStatus state) {
+		if(root.isVisited())
+			return root;
+
+		if(root instanceof DataOp && ((DataOp)root).isWrite() && root.getInput().size() > 0 &&
+			HopRewriteUtils.isMatrixMultiply(root.getInput(0)) && isMMChain(root.getInput()))
+			return optimizeMMChain(root);
+		return root;
+	}
+
+	private static boolean isMMChain(List<Hop> hopList) {
+		boolean ret = true;
+		ret &= (hopList != null && hopList.size() > 0);
+		for(Hop hop : hopList)
+			ret &= isMMChain(hop);
+		return ret;
+	}
+
+	private static boolean isMMChain(Hop hop) {
+		if(HopRewriteUtils.isMatrixMultiply(hop))
+			return isMMChain(hop.getInput());
+		else if(hop instanceof DataOp)
+			return (((DataOp)hop).getOp() == OpOpData.TRANSIENTREAD);
+		return false;
+	}
+
+	public Hop optimizeMMChain(Hop rootHop) {
+		List<Hop> mmChain = new ArrayList<>();
+		List<Hop> mmOperators = new ArrayList<>();
+		Hop mmHop = rootHop.getInput(0);
+		prepareMMChain(mmHop, mmChain, mmOperators);
+
+		Hop newMMHop = sculptMMChain(mmHop, mmChain, mmOperators);
+		if(newMMHop != null) {
+			HopRewriteUtils.replaceChildReference(rootHop, mmHop, newMMHop);
+			mmHop = newMMHop;
+		}
+
+		// print the matrix multiplication chain
+		if(LOG.isTraceEnabled()) {
+			LOG.trace("SparsityOptimizer: Identified and optimized MM Chain: ");
+			for(Hop h : mmChain) {
+				logTraceHop(h, 1);
+			}
+		}
+
+		return rootHop;
+	}
+
+	public void prepareMMChain(Hop root, List<Hop> mmChain, List<Hop> mmOperators) {
+		root.resetVisitStatus();
+		// initialize mmChain with inputs of root hop
+		mmOperators.add(root);
+		for( Hop hi : root.getInput() )
+			mmChain.add(hi);
+
+		// expand each hop to find the entire matrix multiplication chain
+		int counter = 0;
+		while(counter < mmChain.size()) {
+			boolean expandable = true;
+			Hop hop = mmChain.get(counter);
+
+			/*
+			 * Check if mmChain[i] is expandable: 
+			 * 1) It must be MATMULT
+			 * 2) It must not have been visited already
+			 *	(one MATMULT should get expanded only in one chain)
+			 * 3) Its output should not be used in multiple places
+			 *	(either within chain or outside the chain)
+			 */
+			expandable &= HopRewriteUtils.isMatrixMultiply(hop) && !((AggBinaryOp)hop).hasLeftPMInput();
+			expandable &= !hop.isVisited();
+			expandable &= !(hop.getParent().size() > 1);
+
+			hop.setVisited();
+
+			if(!expandable) {
+				counter += 1;
+			}
+			else {
+				List<Hop> tempList = mmChain.get(counter).getInput();
+				// add current operator to mmOperators and its input nodes to mm chain
+				mmOperators.add(mmChain.get(counter));
+				mmChain.set(counter, tempList.get(0));
+				mmChain.add(counter + 1, tempList.get(1));
+			}
+		}
+	}
+
+	// Optimize the matrix multiplication chain by rewriting it
+	// This method can modify the mm chain completely, even re-creating it from scratch is allowed.
+	// NOTE: This method should return the new mm hop when re-creating it from scratch.
+	private Hop sculptMMChain(Hop mmHop, List<Hop> mmChain, List<Hop> mmOperators) {
+		(new RewriteMatrixMultChainOptimizationSparse()).optimizeMMChain(mmHop, mmChain, mmOperators, null);
+		return null;
+	}
+
+	protected void logTraceHop(Hop hop, int level) {
+		if(LOG.isTraceEnabled()) {
+			String offset = Explain.getIdentation(level);
+			LOG.trace(offset + "Hop " + hop.getName() +
+				"(" + hop.getClass().getSimpleName() + ", " + hop.getHopID() + ")" + " " +
+				hop.getDim1() + "x" + hop.getDim2() + " sparsity" + hop.getSparsity());
+		}
+	}
+}

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

@@ -59,10 +59,10 @@ public class RewriteSplitDagSparsityRecompileTest extends AutomatedTestBase {
 	private static final String TEST_DIR = "functions/rewrite/";
 	private static final String TEST_CLASS_DIR =
 		TEST_DIR + RewriteSplitDagSparsityRecompileTest.class.getSimpleName() + "/";
-	private static final String PACKAGE_REWRITE = "org.apache.sysds.hops.rewrite.HopRewriteRule";
 	private static final String PACKAGE_RECOMPILE = "org.apache.sysds.hops.rewrite.StatementBlockRewriteRule";
-	private static Level _oldLevelRewrite;
+	private static final String PACKAGE_SPARSITYOPTIMIZER = "org.apache.sysds.hops.rewrite.SparsityOptimizer";
 	private static Level _oldLevelRecompile;
+	private static Level _oldLevelSparsityOptimizer;
 
 	@Parameterized.Parameter(0)
 	public int rows;
@@ -93,17 +93,17 @@ public void setUp() {
 		addTestConfiguration(TEST_NAME_READ, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME_READ, new String[] {"R"}));
 		addTestConfiguration(TEST_NAME_MMCHAIN, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME_MMCHAIN, new String[] {"R"}));
 		addTestConfiguration(TEST_NAME_MMCHAIN2, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME_MMCHAIN2, new String[] {"R"}));
-		_oldLevelRewrite = Logger.getLogger(PACKAGE_REWRITE).getLevel();
-		Logger.getLogger(PACKAGE_REWRITE).setLevel(Level.TRACE);
 		_oldLevelRecompile = Logger.getLogger(PACKAGE_RECOMPILE).getLevel();
 		Logger.getLogger(PACKAGE_RECOMPILE).setLevel(Level.TRACE);
+		_oldLevelSparsityOptimizer = Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).getLevel();
+		Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).setLevel(Level.TRACE);
 	}
 
 	@Override
 	public void tearDown() {
 		super.tearDown();
-		Logger.getLogger(PACKAGE_REWRITE).setLevel(_oldLevelRewrite);
 		Logger.getLogger(PACKAGE_RECOMPILE).setLevel(_oldLevelRecompile);
+		Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).setLevel(_oldLevelSparsityOptimizer);
 	}
 
 	@Test
@@ -156,7 +156,6 @@ private void testSplitDagSparsityRecompile(String testName, int numSBs) {
 
 			try {
 				DMLConfig dmlConfig = new DMLConfig(getCurConfigFile().getPath());
-				dmlConfig.setTextValue(DMLConfig.SPARSITY_REWRITES, "true");
 				dmlConfig.setTextValue(DMLConfig.SPARSITY_RECOMPILE, "true");
 				overwriteCurrentConfig(dmlConfig);
 			}
@@ -185,13 +184,10 @@ private void testSplitDagSparsityRecompile(String testName, int numSBs) {
 				log_original.stream().map(l -> l.getMessage().toString()).toArray(String[]::new));
 			Assert.assertTrue(log_original_string.contains("RewriteSplitDagSparsityRecompile: Split DAG into " +
 				numSBs + " statement blocks."));
-			Assert.assertTrue(log_original_string.contains("mmchainoptsparse"));
-			Assert.assertTrue(log_original_string.contains("Optimal Sparse MM Chain:"));
 			// check reference log output
 			String log_reference_string = String.join(DELIMITER,
 				log_reference.stream().map(l -> l.getMessage().toString()).toArray(String[]::new));
 			Assert.assertFalse(log_reference_string.contains("RewriteSplitDagSparsityRecompile"));
-			Assert.assertFalse(log_reference_string.contains("mmchainoptsparse"));
 		}
 		finally {
 			OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = oldFlag1;