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 ac3698bdcac..c1751845544 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java @@ -107,7 +107,7 @@ public ProgramRewriter(boolean staticRewrites, boolean dynamicRewrites) if( OptimizerUtils.ALLOW_SPLIT_HOP_DAGS ) _sbRuleSet.add( new RewriteSplitDagUnknownCSVRead() ); //dependency: reblock, merge blocks if(ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_RECOMPILE)) - _sbRuleSet.add(new RewriteSplitDagUnknownNnzRead()); + _sbRuleSet.add(new RewriteSplitDagUnknownNnz()); if( OptimizerUtils.ALLOW_SPLIT_HOP_DAGS && ConfigurationManager.getCompilerConfigFlag(ConfigType.ALLOW_INDIVIDUAL_SB_SPECIFIC_OPS) ) _sbRuleSet.add( new RewriteSplitDagDataDependentOperators() ); //dependency: merge blocks diff --git a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownNnz.java b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownNnz.java new file mode 100644 index 00000000000..2b63b99a6fe --- /dev/null +++ b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownNnz.java @@ -0,0 +1,171 @@ +/* + * 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.Collections; +import java.util.ListIterator; +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; +import org.apache.sysds.parser.DataIdentifier; +import org.apache.sysds.parser.StatementBlock; +import org.apache.sysds.parser.VariableSet; + +/** + * Rule: Split Hop DAG after reads with unknown nnz. This is important to create recompile + * hooks if mtd has unspecified nnz. + */ +public class RewriteSplitDagUnknownNnz extends StatementBlockRewriteRule { + @Override + public boolean createsSplitDag() { + return true; + } + + @Override + public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { + ArrayList ret = new ArrayList<>(); + + // collect all candidate hops in levels based on their data dependencies + ArrayList> candidateLists = new ArrayList>(); + collectNnzComputationCandidates(sb.getHops(), candidateLists); + Collections.reverse(candidateLists); // reverse order of dependency levels + + // split hop dag on demand and if sparsity rewrites are enabled (splitting is pointless w/o rewrite) + if(!candidateLists.isEmpty() && ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_REWRITES)) { + for(ArrayList cand : candidateLists) { + // duplicate sb incl live variable sets + StatementBlock sbSplit = new StatementBlock(); + sbSplit.setDMLProg(sb.getDMLProg()); + sbSplit.setParseInfo(sb); + sbSplit.setLiveIn(new VariableSet()); + sbSplit.setLiveOut(new VariableSet()); + + // move candidate hops to new statement block + // (and replace original persistent read with transient read) + ArrayList sbSplitHops = new ArrayList<>(); + for(Hop cHop : cand) { + // create new transient read + DataOp tRead = HopRewriteUtils.createTransientRead(cHop.getName(), cHop); + tRead.setRequiresRecompile(); + + // replace persistent read with transient read + ArrayList parents = new ArrayList<>(cHop.getParent()); + for(int i = 0; i < parents.size(); i++) { + Hop parent = parents.get(i); + HopRewriteUtils.replaceChildReference(parent, cHop, tRead); + } + + // add reblock sub dag to first statement block + DataOp tWrite = HopRewriteUtils.createTransientWrite(cHop.getName(), cHop); + sbSplitHops.add(tWrite); + + // update live in and out of new statement block (for piggybacking) + DataIdentifier diVar = sb.variablesRead().getVariable(cHop.getName()); + if(diVar != null) { // var read should always exist because persistent read + sbSplit.liveOut().addVariable(cHop.getName(), new DataIdentifier(diVar)); + sb.liveIn().addVariable(cHop.getName(), new DataIdentifier(diVar)); + } + } + + sbSplit.setHops(sbSplitHops); + sbSplit.updateRecompilationFlag(); + ret.add(sbSplit); // statement block with permanent read and transient write + sbSplit.setSplitDag(true); // avoid later merge by other rewrites + } + sb.setSplitDag(true); // avoid later merge by other rewrites + LOG.debug("RewriteSplitDagUnknownNnz: Split DAG into " + (ret.size() + 1) + " statement blocks."); + } + + ret.add(sb); + return ret; + } + + @Override + public List rewriteStatementBlocks(List sbs, ProgramRewriteStatus state) { + return sbs; + } + + private void collectNnzComputationCandidates(ArrayList roots, ArrayList> candidateLists) { + if(roots == null) + return; + Hop.resetVisitStatus(roots); + ListIterator> candListIter = candidateLists.listIterator(); + for(Hop root : roots) { + rCollectNnzComputationCandidates(root, candListIter, null); + } + } + + private ArrayList getNextLevel(ListIterator> liter) { + if(liter.hasNext()) { + return liter.next(); + } + else { + ArrayList newLevel = new ArrayList(); + liter.add(newLevel); + return newLevel; + } + } + + private void rCollectNnzComputationCandidates(Hop hop, ListIterator> candListIter, Hop parentHop) { + if(hop.isVisited()) + return; + + if(!hop.nnzKnown()) { + if(hop instanceof DataOp && ((DataOp) hop).getOp() == OpOpData.PERSISTENTREAD) { + // collect persistent reads with unknown nnz + splitAndProcessChildren(hop, candListIter); + } + else if(parentHop != null && HopRewriteUtils.isMatrixMultiply(parentHop) && + !((AggBinaryOp) parentHop).hasLeftPMInput() && !HopRewriteUtils.isMatrixMultiply(hop)) { + // split before mm chain to ensure correct sparsity information for mm chain sparsity estimation + splitAndProcessChildren(hop, candListIter); + } + else{ + rCollectCandidatesChildren(hop, candListIter); + } + } + else { + rCollectCandidatesChildren(hop, candListIter); + } + + hop.setVisited(); + } + + private void rCollectCandidatesChildren(Hop parentHop, ListIterator> candListIter) { + // process children + if(parentHop.getInput() != null) { + for(Hop c : parentHop.getInput()) + rCollectNnzComputationCandidates(c, candListIter, parentHop); + } + } + + // add the hop to the candidates and proceed with its inputs + private void splitAndProcessChildren(Hop hop, ListIterator> candListIter) { + getNextLevel(candListIter).add(hop); + rCollectCandidatesChildren(hop, candListIter); + candListIter.previous(); // reset level increment after processing children + } +} diff --git a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownNnzRead.java b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownNnzRead.java deleted file mode 100644 index fd39115a982..00000000000 --- a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownNnzRead.java +++ /dev/null @@ -1,147 +0,0 @@ -/* - * 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.sysds.common.Types.OpOpData; -import org.apache.sysds.conf.ConfigurationManager; -import org.apache.sysds.conf.DMLConfig; -import org.apache.sysds.hops.DataOp; -import org.apache.sysds.hops.Hop; -import org.apache.sysds.hops.LiteralOp; -import org.apache.sysds.parser.DataIdentifier; -import org.apache.sysds.parser.StatementBlock; -import org.apache.sysds.parser.VariableSet; - -/** - * Rule: Split Hop DAG after reads with unknown nnz. This is important to create recompile - * hooks if mtd has unspecified nnz. - */ -public class RewriteSplitDagUnknownNnzRead extends StatementBlockRewriteRule { - @Override - public boolean createsSplitDag() { - return true; - } - - @Override - public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { - ArrayList ret = new ArrayList<>(); - - // collect all read hops w/ unknown nnz - ArrayList cand = new ArrayList<>(); - collectReadHopsUnknownNnz(sb.getHops(), cand); - - // split hop dag on demand and if sparsity rewrites are enabled (splitting is pointless w/o rewrite) - if(!cand.isEmpty() && ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_REWRITES)) { - // duplicate sb incl live variable sets - StatementBlock sb1 = new StatementBlock(); - sb1.setDMLProg(sb.getDMLProg()); - sb1.setParseInfo(sb); - sb1.setLiveIn(new VariableSet()); - sb1.setLiveOut(new VariableSet()); - - // move reads incl reblock to new statement block - // (and replace original persistent read with transient read) - ArrayList sb1hops = new ArrayList<>(); - for(Hop reblock : cand) { - // replace reblock inputs to avoid dangling references across dags - // (otherwise, for instance, literal ops are shared across dags) - for(int i = 0; i < reblock.getInput().size(); i++) { - if(reblock.getInput().get(i) instanceof LiteralOp) { - HopRewriteUtils.replaceChildReference(reblock, reblock.getInput().get(i), - new LiteralOp((LiteralOp) reblock.getInput().get(i))); - } - } - - // create new transient read - DataOp tRead = HopRewriteUtils.createTransientRead(reblock.getName(), reblock); - tRead.setRequiresRecompile(); - - // replace reblock with transient read - ArrayList parents = new ArrayList<>(reblock.getParent()); - for(int i = 0; i < parents.size(); i++) { - Hop parent = parents.get(i); - HopRewriteUtils.replaceChildReference(parent, reblock, tRead); - } - - // add reblock sub dag to first statement block - DataOp tWrite = HopRewriteUtils.createTransientWrite(reblock.getName(), reblock); - sb1hops.add(tWrite); - - // update live in and out of new statement block (for piggybacking) - DataIdentifier diVar = sb.variablesRead().getVariable(reblock.getName()); - if(diVar != null) { // var read should always exist because persistent read - sb1.liveOut().addVariable(reblock.getName(), new DataIdentifier(diVar)); - sb.liveIn().addVariable(reblock.getName(), new DataIdentifier(diVar)); - } - } - - sb1.setHops(sb1hops); - sb1.updateRecompilationFlag(); - ret.add(sb1); // statement block with permanent read and transient write - ret.add(sb); // statement block with transient read - sb.setSplitDag(true); // avoid later merge by other rewrites - LOG.debug("Applied splitDagUnknownNnzRead."); - } - // keep original hop dag - else { - ret.add(sb); - } - - return ret; - } - - @Override - public List rewriteStatementBlocks(List sbs, ProgramRewriteStatus state) { - return sbs; - } - - private void collectReadHopsUnknownNnz(ArrayList roots, ArrayList cand) { - if(roots == null) - return; - Hop.resetVisitStatus(roots); - for(Hop root : roots) - collectReadHopsUnknownNnz(root, cand); - } - - private void collectReadHopsUnknownNnz(Hop hop, ArrayList cand) { - if(hop.isVisited()) - return; - - // collect persistent reads with unknown nnz - if(hop instanceof DataOp) { - DataOp dop = (DataOp) hop; - if(dop.getOp() == OpOpData.PERSISTENTREAD && !dop.nnzKnown()) { - cand.add(dop); - } - } - - // process children - if(hop.getInput() != null) { - for(Hop c : hop.getInput()) { - collectReadHopsUnknownNnz(c, cand); - } - } - - hop.setVisited(); - } -} diff --git a/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagUnknownNnzReadTest.java b/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagUnknownNnzTest.java similarity index 94% rename from src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagUnknownNnzReadTest.java rename to src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagUnknownNnzTest.java index 55037e88955..72eaef552ea 100644 --- a/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagUnknownNnzReadTest.java +++ b/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagUnknownNnzTest.java @@ -51,12 +51,12 @@ @RunWith(value = Parameterized.class) @net.jcip.annotations.NotThreadSafe -public class RewriteSplitDagUnknownNnzReadTest extends AutomatedTestBase { +public class RewriteSplitDagUnknownNnzTest extends AutomatedTestBase { private static final String TEST_NAME = "RewriteSplitDagUnknownNnz"; private static final String TEST_DIR = "functions/rewrite/"; private static final String TEST_CLASS_DIR = - TEST_DIR + RewriteSplitDagUnknownNnzReadTest.class.getSimpleName() + "/"; + TEST_DIR + RewriteSplitDagUnknownNnzTest.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; @@ -103,7 +103,7 @@ public void tearDown() { } @Test - public void testSplitDagUnknownNnzRead() { + public void testSplitDagUnknownNnzReadAndMMChain() { boolean oldFlag1 = OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES; boolean oldFlag2 = OptimizerUtils.ALLOW_SUM_PRODUCT_REWRITES; DMLConfig oldDMLConfig = ConfigurationManager.getDMLConfig(); @@ -163,13 +163,13 @@ public void testSplitDagUnknownNnzRead() { // check original log output String log_original_string = String.join(DELIMITER, log_original.stream().map(l -> l.getMessage().toString()).toArray(String[]::new)); - Assert.assertTrue(log_original_string.contains("Applied splitDagUnknownNnzRead.")); + Assert.assertTrue(log_original_string.contains("RewriteSplitDagUnknownNnz: Split DAG into 3 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("splitDagUnknownNnzRead")); + Assert.assertFalse(log_reference_string.contains("RewriteSplitDagUnknownNnz")); Assert.assertFalse(log_reference_string.contains("mmchainoptsparse")); } finally {