diff --git a/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java b/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java index 178a3ecf238..849f2d91e80 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java +++ b/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 diff --git a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java index f0efaaf2e4f..60598d3ef27 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java +++ b/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 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 ret = new ArrayList<>(); // collect all candidate hops in levels based on their data dependencies @@ -158,23 +151,19 @@ private void collectNnzComputationCandidates(ArrayList roots, ArrayList 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 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 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); diff --git a/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java new file mode 100644 index 00000000000..36b84b9cf9f --- /dev/null +++ b/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 rewriteHopDAGs(ArrayList roots, ProgramRewriteStatus state) { + ArrayList ret = new ArrayList(); + 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 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 mmChain = new ArrayList<>(); + List 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 mmChain, List 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 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 mmChain, List 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()); + } + } +} diff --git a/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagSparsityRecompileTest.java b/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagSparsityRecompileTest.java index ac179eb3b9e..5cc57749c2a 100644 --- a/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagSparsityRecompileTest.java +++ b/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; diff --git a/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java b/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java new file mode 100644 index 00000000000..48c7046da3e --- /dev/null +++ b/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java @@ -0,0 +1,192 @@ +/* + * 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 org.apache.log4j.Level; +import org.apache.log4j.Logger; +import org.apache.log4j.spi.LoggingEvent; +import org.apache.sysds.conf.ConfigurationManager; +import org.apache.sysds.conf.DMLConfig; +import org.apache.sysds.hops.OptimizerUtils; +import org.apache.sysds.hops.recompile.Recompiler; +import org.apache.sysds.runtime.matrix.data.MatrixValue; +import org.apache.sysds.test.AutomatedTestBase; +import org.apache.sysds.test.TestConfiguration; +import org.apache.sysds.test.TestUtils; +import org.apache.sysds.test.LoggingUtils; +import org.apache.sysds.test.LoggingUtils.TestAppender; + +import org.junit.Assert; +import org.junit.runners.Parameterized; +import org.junit.Test; +import org.junit.runner.RunWith; + +import java.io.FileNotFoundException; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.stream.DoubleStream; +import java.util.stream.Stream; + +@RunWith(value = Parameterized.class) +@net.jcip.annotations.NotThreadSafe +public class SparsityOptimizerTest extends AutomatedTestBase { + + private static final String TEST_NAME_MMCHAIN = "SparsityOptimizerMMChain"; + private static final String TEST_DIR = "functions/rewrite/"; + private static final String TEST_CLASS_DIR = + TEST_DIR + SparsityOptimizerTest.class.getSimpleName() + "/"; + private static final String PACKAGE_RECOMPILE = "org.apache.sysds.hops.rewrite.StatementBlockRewriteRule"; + private static Level _oldLevelRecompile; + private static final String PACKAGE_SPARSITYOPTIMIZER = "org.apache.sysds.hops.rewrite.SparsityOptimizer"; + private static Level _oldLevelSparsityOptimizer; + + @Parameterized.Parameter(0) + public int rows; + + @Parameterized.Parameter(1) + public int cols; + + @Parameterized.Parameter(2) + public double[] sparsities; + + @Parameterized.Parameter(3) + public double eps; + + @Parameterized.Parameter(4) + public boolean tsmm; + + @Parameterized.Parameters + public static Collection data() { + return Arrays.asList(new Object[][] { + // {rows, cols, sparsities, eps, tsmm}, + {1000, 300, new double[] {0.10d, 0.10d}, Math.pow(10, -10), false}, + {5, 3, new double[] {0.1, 1}, Math.pow(10, -10), true},}); + } + + @Override + public void setUp() { + TestUtils.clearAssertionInformation(); + addTestConfiguration(TEST_NAME_MMCHAIN, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME_MMCHAIN, new String[] {"R"})); + _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_RECOMPILE).setLevel(_oldLevelRecompile); + Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).setLevel(_oldLevelSparsityOptimizer); + } + + @Test + public void testSparsityOptimizationMMChain() { + testSparsityOptimization(TEST_NAME_MMCHAIN); + } + + private void testSparsityOptimization(String testName) { + boolean oldFlag = OptimizerUtils.ALLOW_SPLIT_HOP_DAGS; + DMLConfig oldDMLConfig = ConfigurationManager.getDMLConfig(); + + try { + TestConfiguration config = getTestConfiguration(testName); + loadTestConfiguration(config); + + String HOME = SCRIPT_DIR + TEST_DIR; + fullDMLScriptName = HOME + testName + ".dml"; + programArgs = new String[] {"-explain", "hops", "-stats", + "-args", input("X"), input("Y"), output("R")}; + + OptimizerUtils.ALLOW_SPLIT_HOP_DAGS = true; + + double[][] X = getRandomMatrix(rows, cols, -1, 1, sparsities[0], 7); + double[][] Y = getRandomMatrix(cols, 1, -1, 1, sparsities[1], 3); + long Y_nnz = Stream.of(Y).mapToLong(row -> DoubleStream.of(row).filter(val -> val != 0).count()).sum(); + writeInputMatrixWithMTD("X", X, true); // write matrix characteristics w/ missing nnz + writeInputMatrixWithMTD("Y", Y, Y_nnz, true); + + // run reference test + TestAppender appender = LoggingUtils.overwrite(); // capture log output + runTest(true, false, null, -1); + List log_reference = LoggingUtils.reinsert(appender); // revert the logger to print to stdout + + // get reference result matrix + HashMap referenceResult = readDMLMatrixFromOutputDir("R"); + + try { + DMLConfig dmlConfig = new DMLConfig(getCurConfigFile().getPath()); + dmlConfig.setTextValue(DMLConfig.SPARSITY_RECOMPILE, "true"); + dmlConfig.setTextValue(DMLConfig.SPARSITY_ESTIMATOR, "BASIC_AVG"); + overwriteCurrentConfig(dmlConfig); + } + catch(FileNotFoundException fnfe) { + Assert.fail( + "Could not find DML config file: " + getCurConfigFile().getPath() + " . " + fnfe.getMessage()); + } + catch(IOException ioe) { + Assert.fail("Could not overwrite the DML configuration file. " + ioe.getMessage()); + } + + // run original test + appender = LoggingUtils.overwrite(); // capture log output + runTest(true, false, null, -1); + List log_original = LoggingUtils.reinsert(appender); // revert the logger to print to stdout + + // get resulting matrix + HashMap originalResult = readDMLMatrixFromOutputDir("R"); + + // compare matrices + TestUtils.compareMatrices(originalResult, referenceResult, eps, "Stat-Original", "Stat-Reference"); + + final String DELIMITER = "\n"; + // check original log output + String log_original_string = String.join(DELIMITER, + log_original.stream().map(l -> l.getMessage().toString()).toArray(String[]::new)); + System.out.println(log_original_string); + Assert.assertTrue(log_original_string.contains("RewriteSplitDagSparsityRecompile: Split DAG into")); + Assert.assertTrue(log_original_string.contains("SparsityOptimizer: Identified and optimized 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("SparsityOptimizer")); + } + finally { + OptimizerUtils.ALLOW_SPLIT_HOP_DAGS = oldFlag; + try { + overwriteCurrentConfig(oldDMLConfig); + } + catch(IOException ioe) { + Assert.fail("Unable to restore the previous DML configuration. " + ioe.getMessage()); + } + Recompiler.reinitRecompiler(); + } + } + + private void overwriteCurrentConfig(DMLConfig config) throws IOException { + Files.write(getCurConfigFile().toPath(), config.serializeDMLConfig().getBytes(StandardCharsets.UTF_8)); + } +} diff --git a/src/test/scripts/functions/rewrite/SparsityOptimizerMMChain.dml b/src/test/scripts/functions/rewrite/SparsityOptimizerMMChain.dml new file mode 100644 index 00000000000..2b4308b51f7 --- /dev/null +++ b/src/test/scripts/functions/rewrite/SparsityOptimizerMMChain.dml @@ -0,0 +1,30 @@ +#------------------------------------------------------------- +# +# 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. +# +#------------------------------------------------------------- + + +X = read($1); +Y = read($2); + +S = t(X) %*% X %*% Y; + +R = t(S); + +write(R, $3);