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..178a3ecf238 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java @@ -106,8 +106,9 @@ public ProgramRewriter(boolean staticRewrites, boolean dynamicRewrites) _sbRuleSet.add( new RewriteCompressedReblock() ); // Compression Rewrite 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()); + if( OptimizerUtils.ALLOW_SPLIT_HOP_DAGS && + ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_RECOMPILE) ) + _sbRuleSet.add(new RewriteSplitDagSparsityRecompile()); 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/RewriteSplitDagDataDependentOperators.java b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagDataDependentOperators.java index 3eaed10792b..44f97a7c045 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagDataDependentOperators.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagDataDependentOperators.java @@ -21,7 +21,6 @@ import java.util.ArrayList; import java.util.Arrays; -import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -48,7 +47,6 @@ import org.apache.sysds.parser.DataIdentifier; import org.apache.sysds.parser.StatementBlock; import org.apache.sysds.parser.VariableSet; -import org.apache.sysds.runtime.matrix.data.Pair; /** * Rule: Split Hop DAG after specific data-dependent operators. This is @@ -67,13 +65,8 @@ * split statements blocks again (for consistency). However, currently it is * only used in places (e.g., parfor optimizer) that are not directly affected. */ -public class RewriteSplitDagDataDependentOperators extends StatementBlockRewriteRule +public class RewriteSplitDagDataDependentOperators extends SplitStatementBlockRewriteRule { - @Override - public boolean createsSplitDag() { - return true; - } - @Override public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { @@ -342,140 +335,6 @@ private static Hop getFirstTransientWriteParent( Hop hop ) { return null; } - private void handleReplicatedOperators( ArrayList rootsSB1, ArrayList rootsSB2, VariableSet sb1out, VariableSet sb2in ) - { - //step 1: create probe set SB1 - HashSet probeSet = new HashSet<>(); - Hop.resetVisitStatus(rootsSB1); - for( Hop h : rootsSB1 ) - rAddHopsToProbeSet( h, probeSet ); - - //step 2: probe SB2 operators top-down (collect cut candidates) - HashSet> candSet = new HashSet<>(); - Hop.resetVisitStatus(rootsSB2); - for( Hop h : rootsSB2 ) - rProbeAndAddHopsToCandidateSet(h, probeSet, candSet); - - //step 3: create additional cuts with reuse for common references - HashMap reuseTRead = new HashMap<>(); - for( Pair p : candSet ) { - Hop hop = p.getKey(); - Hop c = p.getValue(); - - DataOp tread = reuseTRead.get(c.getHopID()); - if( tread == null ) { - String varname = createCutVarName(false); - - tread = HopRewriteUtils.createTransientRead(varname, c); - reuseTRead.put(c.getHopID(), tread); - - DataOp twrite = HopRewriteUtils.createTransientWrite(varname, c); - - //update live in and out of new statement block (for piggybacking) - DataIdentifier diVar = new DataIdentifier(varname); - diVar.setDimensions(c.getDim1(), c.getDim2()); - diVar.setBlocksize(c.getBlocksize()); - diVar.setDataType(c.getDataType()); - diVar.setValueType(c.getValueType()); - sb1out.addVariable(varname, new DataIdentifier(diVar)); - sb2in.addVariable(varname, new DataIdentifier(diVar)); - - rootsSB1.add(twrite); - } - - //create additional cut by rewriting both hop dags - int pos = HopRewriteUtils.getChildReferencePos(hop, c); - HopRewriteUtils.removeChildReferenceByPos(hop, c, pos); - HopRewriteUtils.addChildReference(hop, tread, pos); - } - } - - private void rAddHopsToProbeSet( Hop hop, HashSet probeSet ) - { - if( hop.isVisited() ) - return; - - //prevent cuts for no-ops - if( !( (hop instanceof DataOp && !((DataOp)hop).isPersistentReadWrite() ) - || hop instanceof LiteralOp) ) - { - probeSet.add(hop); - } - - if( hop.getInput() != null ) - for( Hop c : hop.getInput() ) - rAddHopsToProbeSet(c, probeSet); - - hop.setVisited(); - } - - /** - * NOTE: candset is a set of parent-child pairs because a parent might have - * multiple references to replicated hops. - * - * @param hop high-level operator - * @param probeSet probe set? - * @param candSet candidate set? - */ - private void rProbeAndAddHopsToCandidateSet( Hop hop, HashSet probeSet, HashSet> candSet ) - { - if( hop.isVisited() ) - return; - - if( hop.getInput() != null ) - for( Hop c : hop.getInput() ) { - //probe for replicated operator, if any child is replicated, keep parent - //for cut between parent-child; otherwise recursively descend. - if( !probeSet.contains(c) ) - rProbeAndAddHopsToCandidateSet(c, probeSet, candSet); - else - { - candSet.add(new Pair<>(hop,c)); - } - } - - hop.setVisited(); - } - - private void collectCandidateChildOperators( ArrayList cand, HashSet candChilds ) - { - Hop.resetVisitStatus(cand); - if( cand != null ) - for( Hop root : cand ) - rCollectCandidateChildOperators(root, cand, candChilds, false); - - // Immediately reset the visit status because candidates might be inner nodes in the DAG. - // Subsequent resets on the root nodes of the DAG would otherwise not necessarily reach - // these nodes which could lead to missing checks on subsequent passes (e.g., when checking - // for replicated operators). - Hop.resetVisitStatus(cand); - } - - private void rCollectCandidateChildOperators( Hop hop, ArrayList cand, HashSet candChilds, boolean collect ) - { - if( hop.isVisited() ) - return; - - //collect operator if necessary - if( collect ) { - candChilds.add(hop); - } - - //activate collection if we passed a candidate - boolean passedFlag = collect; - if( cand.contains(hop) ) { - passedFlag = true; - } - - //process childs recursively - if( hop.getInput()!=null ) { - for( Hop c : hop.getInput() ) - rCollectCandidateChildOperators(c, cand, candChilds, passedFlag); - } - - hop.setVisited(); - } - @Override public List rewriteStatementBlocks(List sbs, ProgramRewriteStatus sate) { return sbs; diff --git a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java new file mode 100644 index 00000000000..c3237677ce9 --- /dev/null +++ b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java @@ -0,0 +1,212 @@ +/* + * 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.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; +import org.apache.sysds.hops.recompile.Recompiler; +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 RewriteSplitDagSparsityRecompile extends SplitStatementBlockRewriteRule { + @Override + public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { + return rRewriteStatementBlock(sb, state); + } + + 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 + ArrayList cand = new ArrayList(); + collectNnzComputationCandidates(sb.getHops(), cand); + + // split hop dag on demand + if(!cand.isEmpty()) { + //collect child operators of candidates (to prevent rewrite anomalies) + //For example, B = rmEmpty(A), C = rmEmpty(B), op(C, nrow(B)) needs to + //preserve the link rmEmpty(B) in the split-off DAG. Therefore, the + //candidates themselves need to be part of this anomaly filter as well. + HashSet candChilds = new HashSet<>(); + collectCandidateChildOperators(cand, candChilds); + candChilds.addAll(cand); + + // 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 + String varname = createCutVarName(false); + DataOp tRead = HopRewriteUtils.createTransientRead(varname, cHop); + tRead.setRequiresRecompile(); + + // replace persistent read with transient read + ArrayList parents = new ArrayList<>(cHop.getParent()); + for(int i = 0; i < parents.size(); i++) { + // prevent concurrent modification by index access + Hop parent = parents.get(i); + if(!candChilds.contains(parent)) // anomaly filter + HopRewriteUtils.replaceChildReference(parent, cHop, tRead); + } + + // add reblock sub dag to first statement block + DataOp tWrite = HopRewriteUtils.createTransientWrite(varname, cHop); + sbSplitHops.add(tWrite); + + // update live in and out of new statement block (for piggybacking) + DataIdentifier diVar = new DataIdentifier(varname); + diVar.setDimensions(cHop.getDim1(), cHop.getDim2()); + diVar.setBlocksize(cHop.getBlocksize()); + diVar.setDataType(cHop.getDataType()); + diVar.setValueType(cHop.getValueType()); + sbSplit.liveOut().addVariable(varname, new DataIdentifier(diVar)); + sb.liveIn().addVariable(varname, new DataIdentifier(diVar)); + sb.variablesRead().addVariable(varname, new DataIdentifier(diVar)); + } + + // ensure disjoint operators across DAGs (prevent replicated operations) + handleReplicatedOperators(sbSplitHops, sb.getHops(), sbSplit.liveOut(), sb.liveIn()); + + sbSplit.setHops(Recompiler.deepCopyHopsDag(sbSplitHops)); + sbSplit.updateRecompilationFlag(); + sbSplit.setSplitDag(true); // avoid later merge by other rewrites + + //recursive application of rewrite rule + List tmp = rRewriteStatementBlock(sbSplit, state); + + ret.addAll(tmp); // statement block with permanent read and transient write + sb.setSplitDag(true); // avoid later merge by other rewrites + LOG.debug("RewriteSplitDagSparsityRecompile: 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 cand) { + if(roots == null) + return; + Hop.resetVisitStatus(roots); + for(Hop root : roots) { + rCollectNnzComputationCandidates(root, cand, null); + } + } + + public static boolean anyParentIsMatrixMultiply(Hop hop) { + return hop.getParent().stream().anyMatch(parentHop -> HopRewriteUtils.isMatrixMultiply(parentHop)); + } + + public static boolean allParentsAreMatrixMultiply(Hop hop) { + return hop.getParent().stream().allMatch(parentHop -> HopRewriteUtils.isMatrixMultiply(parentHop)); + } + + public 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) { + // hop is part of mm chain with at least two matrix multplications + return HopRewriteUtils.isMatrixMultiply(hop) + && (anyParentIsMatrixMultiply(hop) || anyChildIsMatrixMultiply(hop)); + } + + public static boolean anyParentIsMMChain(Hop hop) { + List parentHops = hop.getParent(); + if(parentHops == null) + return false; + return parentHops.stream().anyMatch(pHop -> isMMChain(pHop)); + } + + private void rCollectNnzComputationCandidates(Hop hop, ArrayList cand, Hop parentHop) { + if(hop.isVisited()) + return; + // has already been split off + boolean noSplitRequired = HopRewriteUtils.hasOnlyWriteParents(hop, true, false); + + if(!noSplitRequired) { + if(hop instanceof DataOp && ((DataOp) hop).getOp() == OpOpData.PERSISTENTREAD && !hop.nnzKnown()) { + // split persistent reads with unknown nnz to obtain their nnz before continuing rewrites + cand.add(hop); + } + else if(parentHop != null && !HopRewriteUtils.isMatrixMultiply(hop) + && !(hop instanceof DataOp && ((DataOp) hop).getOp() == OpOpData.TRANSIENTREAD) + && anyParentIsMMChain(hop)) { // split only mm chains (i.e., at least two matrix multiplications) + // split before mm chain to ensure correct sparsity information for mm chain sparsity estimation + cand.add(hop); + } + else if((!HopRewriteUtils.isMatrixMultiply(parentHop) || !allParentsAreMatrixMultiply(hop)) + && isMMChain(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); + } + else { + rCollectNnzComputationCandidatesChildren(hop, cand); + } + } + else { + rCollectNnzComputationCandidatesChildren(hop, cand); + } + + hop.setVisited(); + } + + private void rCollectNnzComputationCandidatesChildren(Hop parentHop, ArrayList cand) { + // process children + if(parentHop.getInput() != null) { + for(Hop c : parentHop.getInput()) + rCollectNnzComputationCandidates(c, cand, parentHop); + } + } +} diff --git a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownCSVRead.java b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownCSVRead.java index a6cf158ad84..33b9861122d 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownCSVRead.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagUnknownCSVRead.java @@ -42,13 +42,8 @@ * mtd has unknown size (which can only happen for CSV). * */ -public class RewriteSplitDagUnknownCSVRead extends StatementBlockRewriteRule +public class RewriteSplitDagUnknownCSVRead extends SplitStatementBlockRewriteRule { - @Override - public boolean createsSplitDag() { - return true; - } - @Override public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { 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/main/java/org/apache/sysds/hops/rewrite/SplitStatementBlockRewriteRule.java b/src/main/java/org/apache/sysds/hops/rewrite/SplitStatementBlockRewriteRule.java new file mode 100644 index 00000000000..eeda6bd704a --- /dev/null +++ b/src/main/java/org/apache/sysds/hops/rewrite/SplitStatementBlockRewriteRule.java @@ -0,0 +1,168 @@ +/* + * 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.HashMap; +import java.util.HashSet; + +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.VariableSet; +import org.apache.sysds.runtime.matrix.data.Pair; + +public abstract class SplitStatementBlockRewriteRule extends StatementBlockRewriteRule { + @Override + public boolean createsSplitDag() { + return true; + } + + protected void handleReplicatedOperators(ArrayList rootsSbSplit, ArrayList rootsSbOrigin, + VariableSet sbSplitOut, VariableSet sbOriginIn) { + // step 1: create probe set SB1 + HashSet probeSet = new HashSet<>(); + Hop.resetVisitStatus(rootsSbSplit); + for(Hop h : rootsSbSplit) { + rAddHopsToProbeSet(h, probeSet); + } + + // step 2: probe SB2 operators top-down (collect cut candidates) + HashSet> candSet = new HashSet<>(); + Hop.resetVisitStatus(rootsSbOrigin); + for(Hop h : rootsSbOrigin) { + rProbeAndAddHopsToCandidateSet(h, probeSet, candSet); + } + + // step 3: create additional cuts with reuse for common references + HashMap reuseTRead = new HashMap<>(); + for(Pair p : candSet) { + Hop hop = p.getKey(); + Hop c = p.getValue(); + + DataOp tRead = reuseTRead.get(c.getHopID()); + if(tRead == null) { + String varname = createCutVarName(false); + + tRead = HopRewriteUtils.createTransientRead(varname, c); + reuseTRead.put(c.getHopID(), tRead); + + DataOp twrite = HopRewriteUtils.createTransientWrite(varname, c); + + // update live in and out of new statement block (for piggybacking) + DataIdentifier diVar = new DataIdentifier(varname); + diVar.setDimensions(c.getDim1(), c.getDim2()); + diVar.setBlocksize(c.getBlocksize()); + diVar.setDataType(c.getDataType()); + diVar.setValueType(c.getValueType()); + sbSplitOut.addVariable(varname, new DataIdentifier(diVar)); + sbOriginIn.addVariable(varname, new DataIdentifier(diVar)); + + rootsSbSplit.add(twrite); + } + + // create additional cut by rewriting both hop dags + int pos = HopRewriteUtils.getChildReferencePos(hop, c); + HopRewriteUtils.removeChildReferenceByPos(hop, c, pos); + HopRewriteUtils.addChildReference(hop, tRead, pos); + } + } + + private void rAddHopsToProbeSet(Hop hop, HashSet probeSet) { + if(hop.isVisited()) + return; + + // prevent cuts for no-ops + if(!(hop instanceof DataOp && !((DataOp)hop).isPersistentReadWrite()) && !(hop instanceof LiteralOp)) + probeSet.add(hop); + + if(hop.getInput() != null) { + for(Hop c : hop.getInput()) { + rAddHopsToProbeSet(c, probeSet); + } + } + + hop.setVisited(); + } + + /** + * NOTE: candset is a set of parent-child pairs because a parent might have + * multiple references to replicated hops. + * + * @param hop high-level operator + * @param probeSet probe set? + * @param candSet candidate set? + */ + private void rProbeAndAddHopsToCandidateSet(Hop hop, HashSet probeSet, HashSet> candSet) { + if(hop.isVisited()) + return; + + if(hop.getInput() != null) + for(Hop c : hop.getInput()) { + // probe for replicated operator, if any child is replicated, keep parent + // for cut between parent-child; otherwise recursively descend. + if(!probeSet.contains(c)) + rProbeAndAddHopsToCandidateSet(c, probeSet, candSet); + else + candSet.add(new Pair<>(hop,c)); + } + + hop.setVisited(); + } + + protected void collectCandidateChildOperators(ArrayList cand, HashSet candChilds) { + Hop.resetVisitStatus(cand); + if(cand != null) { + for(Hop root : cand) { + rCollectCandidateChildOperators(root, cand, candChilds, false); + } + } + + // Immediately reset the visit status because candidates might be inner nodes in the DAG. + // Subsequent resets on the root nodes of the DAG would otherwise not necessarily reach + // these nodes which could lead to missing checks on subsequent passes (e.g., when checking + // for replicated operators). + Hop.resetVisitStatus(cand); + } + + private void rCollectCandidateChildOperators(Hop hop, ArrayList cand, HashSet candChilds, boolean collect) { + if(hop.isVisited()) + return; + + // collect operator if necessary + if(collect) + candChilds.add(hop); + + // activate collection if we passed a hop from the candidates + boolean passedFlag = collect; + if(cand.contains(hop)) + passedFlag = true; + + // process childs recursively + if(hop.getInput() != null) { + for(Hop childHop : hop.getInput()) { + rCollectCandidateChildOperators(childHop, cand, candChilds, passedFlag); + } + } + + 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/RewriteSplitDagSparsityRecompileTest.java similarity index 82% rename from src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagUnknownNnzReadTest.java rename to src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagSparsityRecompileTest.java index 55037e88955..1802ef54660 100644 --- a/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagUnknownNnzReadTest.java +++ b/src/test/java/org/apache/sysds/test/functions/rewrite/RewriteSplitDagSparsityRecompileTest.java @@ -51,12 +51,13 @@ @RunWith(value = Parameterized.class) @net.jcip.annotations.NotThreadSafe -public class RewriteSplitDagUnknownNnzReadTest extends AutomatedTestBase { +public class RewriteSplitDagSparsityRecompileTest extends AutomatedTestBase { - private static final String TEST_NAME = "RewriteSplitDagUnknownNnz"; + private static final String TEST_NAME_READ = "RewriteSplitDagSparsityRecompileRead"; + private static final String TEST_NAME_MMCHAIN = "RewriteSplitDagSparsityRecompileMMChain"; private static final String TEST_DIR = "functions/rewrite/"; private static final String TEST_CLASS_DIR = - TEST_DIR + RewriteSplitDagUnknownNnzReadTest.class.getSimpleName() + "/"; + 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; @@ -80,7 +81,7 @@ public class RewriteSplitDagUnknownNnzReadTest extends AutomatedTestBase { @Parameterized.Parameters public static Collection data() { return Arrays.asList(new Object[][] { - // {rows, cols, sparsities, eps}, + // {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},}); } @@ -88,7 +89,8 @@ public static Collection data() { @Override public void setUp() { TestUtils.clearAssertionInformation(); - addTestConfiguration(TEST_NAME, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME, new String[] {"R"})); + 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"})); _oldLevelRewrite = Logger.getLogger(PACKAGE_REWRITE).getLevel(); Logger.getLogger(PACKAGE_REWRITE).setLevel(Level.TRACE); _oldLevelRecompile = Logger.getLogger(PACKAGE_RECOMPILE).getLevel(); @@ -103,22 +105,33 @@ public void tearDown() { } @Test - public void testSplitDagUnknownNnzRead() { + public void testSplitDagSparsityRecompileRead() { + testSplitDagSparsityRecompile(TEST_NAME_READ, 4); + } + + @Test + public void testSplitDagSparsityRecompileMMChain() { + testSplitDagSparsityRecompile(TEST_NAME_MMCHAIN, 4); + } + + private void testSplitDagSparsityRecompile(String testName, int numSBs) { boolean oldFlag1 = OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES; boolean oldFlag2 = OptimizerUtils.ALLOW_SUM_PRODUCT_REWRITES; + boolean oldFlag3 = OptimizerUtils.ALLOW_SPLIT_HOP_DAGS; DMLConfig oldDMLConfig = ConfigurationManager.getDMLConfig(); try { - TestConfiguration config = getTestConfiguration(TEST_NAME); + TestConfiguration config = getTestConfiguration(testName); loadTestConfiguration(config); String HOME = SCRIPT_DIR + TEST_DIR; - fullDMLScriptName = HOME + TEST_NAME + ".dml"; + fullDMLScriptName = HOME + testName + ".dml"; programArgs = new String[] {"-explain", "hops", "-stats", "-args", input("X"), input("Y"), output("R")}; OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = true; OptimizerUtils.ALLOW_SUM_PRODUCT_REWRITES = true; + 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); @@ -163,18 +176,20 @@ 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("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("splitDagUnknownNnzRead")); + Assert.assertFalse(log_reference_string.contains("RewriteSplitDagSparsityRecompile")); Assert.assertFalse(log_reference_string.contains("mmchainoptsparse")); } finally { OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = oldFlag1; OptimizerUtils.ALLOW_SUM_PRODUCT_REWRITES = oldFlag2; + OptimizerUtils.ALLOW_SPLIT_HOP_DAGS = oldFlag3; try { overwriteCurrentConfig(oldDMLConfig); } diff --git a/src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileMMChain.dml b/src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileMMChain.dml new file mode 100644 index 00000000000..aa305406602 --- /dev/null +++ b/src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileMMChain.dml @@ -0,0 +1,27 @@ +#------------------------------------------------------------- +# +# 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); +R = t(X) %*% X %*% t(t(X) %*% X) %*% Y; + +write(R, $3); diff --git a/src/test/scripts/functions/rewrite/RewriteSplitDagUnknownNnz.dml b/src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileRead.dml similarity index 99% rename from src/test/scripts/functions/rewrite/RewriteSplitDagUnknownNnz.dml rename to src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileRead.dml index 998334b611c..498da9ec519 100644 --- a/src/test/scripts/functions/rewrite/RewriteSplitDagUnknownNnz.dml +++ b/src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileRead.dml @@ -22,7 +22,6 @@ X = read($1); Y = read($2); - R = t(X) %*% X %*% Y; write(R, $3);