[SYSTEMML-1828,1832] New rewrite for merging statement block sequences

This patch introduces a new statement block rewrite for merging DAGs of
subsequent last-level statement blocks. After constant folding and the
removal of unnecessary branches, we often end up with such sequences of
statement blocks. Since many rewrites and operator fusion work on the
granularity of individual DAGs, these unnecessary DAG cuts cause missed
optimization opportunities, especially in the context of operator fusion
(i.e., codegen). We now merge such sequences in awareness of rewrites
that explicitly split DAGs (to create recompilation points).

Apart from the new merge rewrite, this patch also fixes the IPA rewrite
pass that applies static rewrites per IPA round. The repeated
application of the statement block rewrite for injecting spark
checkpoints for variables used read-only in loops introduced redundant
statement blocks and checkpoints. The IPA rewrite pass now explicitly
excludes this rewrite.

Additionally, this patch also modifies the related tests to use
'while(FALSE){}' instead of 'if(1==1){}' as a DAG cut, and fixes some
minor compilation issues that showed up due to the increases
optimization scope.

Overall, there are many scripts and patterns that benefit from these
changes. For example, on 1 epoch of lenet w/ codegen, this patch
improved end-to-end performance from 328s to 297s due to increased
fusion opportunities and fewer compiled spark instructions (70 vs 82).

Author: mboehm7

docs/dml-language-reference.md

@@ -1782,7 +1782,7 @@ The following DML utilizes the `transformencode()` function.
     jspec = read("/user/ml/homes.tfspec_recode2.json", data_type="scalar", value_type="string");
     [X, M] = transformencode(target=F1, spec=jspec);
     print(toString(X));
-    if(1==1){}
+    while(FALSE){}
     print(toString(M));
 
 The transformed matrix X and output M are as follows.

projects/breast_cancer/MachineLearning.ipynb

@@ -683,7 +683,7 @@
    "source": [
     "# script = \"\"\"\n",
     "# f = function(matrix[double] X) return(matrix[double] Y) {\n",
-    "#   if (1==1) {}\n",
+    "#   while(FALSE){}\n",
     "#   a = as.scalar(rand(rows=1, cols=1))\n",
     "#   Y = X * a\n",
     "# }\n",

scripts/nn/test/test.dml

@@ -599,7 +599,7 @@ max_pool2d = function() {
   for (padh in 0:3) {
     for (padw in 0:3) {
       print(" - Testing w/ padh="+padh+" & padw="+padw+".")
-      #if (1==1) {}  # force correct printing
+      #while(FALSE){}  # force correct printing
       #print("   - Testing forward")
       [out, Hout, Wout] = max_pool2d::forward(X, C, Hin, Win, Hf, Wf, stride, stride, padh, padw)
       [out_simple, Hout_simple, Wout_simple] = max_pool2d_simple::forward(X, C, Hin, Win, Hf, Wf,

src/main/java/org/apache/sysml/hops/ipa/IPAPassApplyStaticHopRewrites.java

@@ -22,6 +22,7 @@
 
 import org.apache.sysml.hops.HopsException;
 import org.apache.sysml.hops.rewrite.ProgramRewriter;
+import org.apache.sysml.hops.rewrite.RewriteInjectSparkLoopCheckpointing;
 import org.apache.sysml.parser.DMLProgram;
 import org.apache.sysml.parser.LanguageException;
 
@@ -43,7 +44,11 @@ public void rewriteProgram( DMLProgram prog, FunctionCallGraph fgraph, FunctionC
 		throws HopsException
 	{
 		try {
+			//construct rewriter w/o checkpoint injection to avoid redundancy
 			ProgramRewriter rewriter = new ProgramRewriter(true, false);
+			rewriter.removeStatementBlockRewrite(RewriteInjectSparkLoopCheckpointing.class);
+			
+			//rewrite program hop dags and statement blocks
 			rewriter.rewriteProgramHopDAGs(prog);
 		} 
 		catch (LanguageException ex) {

src/main/java/org/apache/sysml/hops/rewrite/HopRewriteUtils.java

@@ -622,6 +622,14 @@ public static TernaryOp createTernaryOp(Hop mleft, Hop smid, Hop mright, OpOp3 o
 		return ternOp;
 	}
 
+	public static DataOp createDataOp(String name, Hop input, DataOpTypes type) {
+		DataOp dop = new DataOp(name, input.getDataType(), input.getValueType(), input, type, null);
+		dop.setOutputBlocksizes(input.getRowsInBlock(), input.getColsInBlock());
+		copyLineNumbers(input, dop);
+		dop.refreshSizeInformation();
+		return dop;
+	}
+	
 	public static void setOutputParameters( Hop hop, long rlen, long clen, long brlen, long bclen, long nnz ) {
 		hop.setDim1( rlen );
 		hop.setDim2( clen );
@@ -846,6 +854,11 @@ private static boolean rContainsInput(Hop current, Hop probe, HashSet<Long> memo
 		return ret;
 	}
 
+	public static boolean isData(Hop hop, DataOpTypes type) {
+		return hop instanceof DataOp
+			&& ((DataOp)hop).getDataOpType()==type;
+	}
+	
 	public static boolean isBinaryMatrixColVectorOperation(Hop hop) {
 		return hop instanceof BinaryOp 
 			&& hop.getInput().get(0).getDataType().isMatrix() && hop.getInput().get(1).getDataType().isMatrix()

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

@@ -20,6 +20,7 @@
 package org.apache.sysml.hops.rewrite;
 
 import java.util.ArrayList;
+import java.util.List;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -49,7 +50,7 @@
  * program. 
  * 
  */
-public class ProgramRewriter 
+public class ProgramRewriter
 {
 	private static final Log LOG = LogFactory.getLog(ProgramRewriter.class.getName());
 
@@ -60,17 +61,15 @@ public class ProgramRewriter
 	private ArrayList<HopRewriteRule> _dagRuleSet = null;
 	private ArrayList<StatementBlockRewriteRule> _sbRuleSet = null;
 
-	static{
+	static {
 		// for internal debugging only
 		if( LDEBUG ) {
 			Logger.getLogger("org.apache.sysml.hops.rewrite")
 				  .setLevel((Level) Level.DEBUG);
 		}
-		
 	}
 
-	public ProgramRewriter()
-	{
+	public ProgramRewriter() {
 		// by default which is used during initial compile 
 		// apply all (static and dynamic) rewrites
 		this( true, true );
@@ -107,12 +106,14 @@ public ProgramRewriter( boolean staticRewrites, boolean dynamicRewrites )
 			_dagRuleSet.add( new RewriteInjectSparkPReadCheckpointing()          ); //dependency: reblock
 
 			//add statement block rewrite rules
- 			if( OptimizerUtils.ALLOW_BRANCH_REMOVAL )			
+ 			if( OptimizerUtils.ALLOW_BRANCH_REMOVAL ) {
 				_sbRuleSet.add(  new RewriteRemoveUnnecessaryBranches()          ); //dependency: constant folding		
+				_sbRuleSet.add(  new RewriteMergeBlockSequence()                 ); //dependency: remove branches
+ 			}
  			if( OptimizerUtils.ALLOW_SPLIT_HOP_DAGS )
- 				_sbRuleSet.add(  new RewriteSplitDagUnknownCSVRead()             ); //dependency: reblock	
+ 				_sbRuleSet.add(  new RewriteSplitDagUnknownCSVRead()             ); //dependency: reblock, merge blocks	
  			if( ConfigurationManager.getCompilerConfigFlag(ConfigType.ALLOW_INDIVIDUAL_SB_SPECIFIC_OPS) )
- 				_sbRuleSet.add(  new RewriteSplitDagDataDependentOperators()     );
+ 				_sbRuleSet.add(  new RewriteSplitDagDataDependentOperators()     ); //dependency: merge blocks
  			if( OptimizerUtils.ALLOW_AUTO_VECTORIZATION )
 				_sbRuleSet.add(  new RewriteForLoopVectorization()               ); //dependency: reblock (reblockop)
  			_sbRuleSet.add( new RewriteInjectSparkLoopCheckpointing(true)        ); //dependency: reblock (blocksizes)
@@ -146,8 +147,7 @@ public ProgramRewriter( boolean staticRewrites, boolean dynamicRewrites )
 	 * 
 	 * @param rewrites the HOP rewrite rules
 	 */
-	public ProgramRewriter( HopRewriteRule... rewrites )
-	{
+	public ProgramRewriter( HopRewriteRule... rewrites ) {
 		//initialize HOP DAG rewrite ruleSet (with fixed rewrite order)
 		_dagRuleSet = new ArrayList<HopRewriteRule>();
 		for( HopRewriteRule rewrite : rewrites )
@@ -161,8 +161,7 @@ public ProgramRewriter( HopRewriteRule... rewrites )
 	 * 
 	 * @param rewrites the statement block rewrite rules
 	 */
-	public ProgramRewriter( StatementBlockRewriteRule... rewrites )
-	{
+	public ProgramRewriter( StatementBlockRewriteRule... rewrites ) {
 		//initialize HOP DAG rewrite ruleSet (with fixed rewrite order)
 		_dagRuleSet = new ArrayList<HopRewriteRule>();
 
@@ -177,8 +176,7 @@ public ProgramRewriter( StatementBlockRewriteRule... rewrites )
 	 * @param hRewrites HOP rewrite rules
 	 * @param sbRewrites statement block rewrite rules
 	 */
-	public ProgramRewriter( ArrayList<HopRewriteRule> hRewrites, ArrayList<StatementBlockRewriteRule> sbRewrites )
-	{
+	public ProgramRewriter(ArrayList<HopRewriteRule> hRewrites, ArrayList<StatementBlockRewriteRule> sbRewrites) {
 		//initialize HOP DAG rewrite ruleSet (with fixed rewrite order)
 		_dagRuleSet = new ArrayList<HopRewriteRule>();
 		_dagRuleSet.addAll( hRewrites );
@@ -187,6 +185,14 @@ public ProgramRewriter( ArrayList<HopRewriteRule> hRewrites, ArrayList<Statement
 		_sbRuleSet.addAll( sbRewrites );
 	}
 
+	public void removeHopRewrite(Class<? extends HopRewriteRule> clazz) {
+		_dagRuleSet.removeIf(r -> r.getClass().equals(clazz));
+	}
+	
+	public void removeStatementBlockRewrite(Class<? extends StatementBlockRewriteRule> clazz) {
+		_sbRuleSet.removeIf(r -> r.getClass().equals(clazz));
+	}
+	
 	public ProgramRewriteStatus rewriteProgramHopDAGs(DMLProgram dmlp) 
 		throws LanguageException, HopsException
 	{	
@@ -301,21 +307,23 @@ public Hop rewriteHopDAG(Hop root, ProgramRewriteStatus state)
 		return root;
 	}
 
-	public ArrayList<StatementBlock> rewriteStatementBlocks( ArrayList<StatementBlock> sbs, ProgramRewriteStatus state ) 
+	public ArrayList<StatementBlock> rewriteStatementBlocks( ArrayList<StatementBlock> sbs, ProgramRewriteStatus status ) 
 		throws HopsException
 	{
 		//ensure robustness for calls from outside
-		if( state == null )
-			state = new ProgramRewriteStatus();
-				
+		if( status == null )
+			status = new ProgramRewriteStatus();
 
-		ArrayList<StatementBlock> tmp = new ArrayList<StatementBlock>();
+		//apply rewrite rules to list of statement blocks
+		List<StatementBlock> sbList = sbs; 
+		for( StatementBlockRewriteRule r : _sbRuleSet ) {
+			sbList = r.rewriteStatementBlocks(sbList, status);
+		}
 
 		//rewrite statement blocks (with potential expansion)
-		for( StatementBlock sb : sbs )
-			tmp.addAll( rewriteStatementBlock(sb, state) );
-		
-		//copy results into original collection
+		ArrayList<StatementBlock> tmp = new ArrayList<StatementBlock>();
+		for( StatementBlock sb : sbList )
+			tmp.addAll( rewriteStatementBlock(sb, status) );
 		sbs.clear();
 		sbs.addAll( tmp );
 
@@ -362,9 +370,8 @@ else if (sb instanceof ForStatementBlock) //incl parfor
 			status.setInParforContext(prestatus);
 		}
 
-		//apply rewrite rules
-		for( StatementBlockRewriteRule r : _sbRuleSet )
-		{
+		//apply rewrite rules to individual statement blocks
+		for( StatementBlockRewriteRule r : _sbRuleSet ) {
 			ArrayList<StatementBlock> tmp = new ArrayList<StatementBlock>();			
 			for( StatementBlock sbc : ret )
 				tmp.addAll( r.rewriteStatementBlock(sbc, status) );

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

@@ -1647,11 +1647,9 @@ private Hop simplifyOuterSeqExpand(Hop parent, Hop hi, int pos)
 				inputargs.put("cast", new LiteralOp(false));
 
 				//create new hop
-				ParameterizedBuiltinOp pbop = new ParameterizedBuiltinOp("tmp", DataType.MATRIX, ValueType.DOUBLE, 
-						ParamBuiltinOp.REXPAND, inputargs);
-				pbop.setOutputBlocksizes(hi.getRowsInBlock(), hi.getColsInBlock());
-				pbop.refreshSizeInformation();
-		
+				ParameterizedBuiltinOp pbop = HopRewriteUtils
+					.createParameterizedBuiltinOp(trgt, inputargs, ParamBuiltinOp.REXPAND);
+				
 				//relink new hop into original position
 				HopRewriteUtils.replaceChildReference(parent, hi, pbop, pos);
 				hi = pbop;

src/main/java/org/apache/sysml/hops/rewrite/RewriteForLoopVectorization.java

@@ -19,7 +19,8 @@
 
 package org.apache.sysml.hops.rewrite;
 
-import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 
 import org.apache.sysml.hops.AggUnaryOp;
 import org.apache.sysml.hops.BinaryOp;
@@ -49,17 +50,14 @@
  */
 public class RewriteForLoopVectorization extends StatementBlockRewriteRule
 {
-
 	private static final OpOp2[] MAP_SCALAR_AGGREGATE_SOURCE_OPS = new OpOp2[]{OpOp2.PLUS, OpOp2.MULT, OpOp2.MIN, OpOp2.MAX};
 	private static final AggOp[] MAP_SCALAR_AGGREGATE_TARGET_OPS = new AggOp[]{AggOp.SUM,  AggOp.PROD, AggOp.MIN, AggOp.MAX};
 
 
 	@Override
-	public ArrayList<StatementBlock> rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state)
+	public List<StatementBlock> rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state)
 		throws HopsException 
 	{
-		ArrayList<StatementBlock> ret = new ArrayList<StatementBlock>();
-		
 		if( sb instanceof ForStatementBlock )
 		{
 			ForStatementBlock fsb = (ForStatementBlock) sb;
@@ -96,9 +94,13 @@ public ArrayList<StatementBlock> rewriteStatementBlock(StatementBlock sb, Progra
 
 		//if no rewrite applied sb is the original for loop otherwise a last level statement block
 		//that includes the equivalent vectorized operations.
-		ret.add( sb );
-		
-		return ret;
+		return Arrays.asList(sb);
+	}
+	
+	@Override
+	public List<StatementBlock> rewriteStatementBlocks(List<StatementBlock> sbs, 
+			ProgramRewriteStatus sate) throws HopsException {
+		return sbs;
 	}
 
 	private StatementBlock vectorizeScalarAggregate( StatementBlock sb, StatementBlock csb, Hop from, Hop to, Hop increment, String itervar ) 

src/main/java/org/apache/sysml/hops/rewrite/RewriteInjectSparkLoopCheckpointing.java

@@ -20,6 +20,8 @@
 package org.apache.sysml.hops.rewrite;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 
 import org.apache.sysml.hops.DataOp;
 import org.apache.sysml.hops.Hop;
@@ -47,28 +49,25 @@ public class RewriteInjectSparkLoopCheckpointing extends StatementBlockRewriteRu
 {
 	private boolean _checkCtx = false;
 
-	public RewriteInjectSparkLoopCheckpointing(boolean checkParForContext)
-	{
+	public RewriteInjectSparkLoopCheckpointing(boolean checkParForContext) {
 		_checkCtx = checkParForContext;
 	}
 
 	@Override
-	public ArrayList<StatementBlock> rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus status)
+	public List<StatementBlock> rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus status)
 		throws HopsException 
 	{
-		ArrayList<StatementBlock> ret = new ArrayList<StatementBlock>();
-		
-		if( !OptimizerUtils.isSparkExecutionMode() ) 
-		{
-			ret.add(sb); // nothing to do here
-			return ret; //return original statement block
+		if( !OptimizerUtils.isSparkExecutionMode() ) {
+			// nothing to do here, return original statement block
+			return Arrays.asList(sb);
 		}
 
 		//1) We currently add checkpoint operations without information about the global program structure,
 		//this assumes that redundant checkpointing is prevented at runtime level (instruction-level)
 		//2) Also, we do not take size information into account right now. This means that all candidates
 		//are checkpointed even if they are only used by CP operations.
 
+		ArrayList<StatementBlock> ret = new ArrayList<>();
 		int blocksize = status.getBlocksize(); //block size set by reblock rewrite
 
 		//apply rewrite for while, for, and parfor (the decision for parfor loop bodies is deferred until parfor
@@ -101,7 +100,7 @@ public ArrayList<StatementBlock> rewriteStatementBlock(StatementBlock sb, Progra
 					long dim2 = (dat instanceof IndexedIdentifier) ? ((IndexedIdentifier)dat).getOrigDim2() : dat.getDim2();
 					DataOp tread = new DataOp(var, DataType.MATRIX, ValueType.DOUBLE, DataOpTypes.TRANSIENTREAD, 
 							            dat.getFilename(), dim1, dim2, dat.getNnz(), blocksize, blocksize);
-					tread.setRequiresCheckpoint( true );
+					tread.setRequiresCheckpoint(true);
 					DataOp twrite = new DataOp(var, DataType.MATRIX, ValueType.DOUBLE, tread, DataOpTypes.TRANSIENTWRITE, null);
 					HopRewriteUtils.setOutputParameters(twrite, dim1, dim2, blocksize, blocksize, dat.getNnz());					
 					hops.add(twrite);
@@ -111,6 +110,7 @@ public ArrayList<StatementBlock> rewriteStatementBlock(StatementBlock sb, Progra
 				sb0.set_hops(hops);
 				sb0.setLiveIn(livein);
 				sb0.setLiveOut(liveout);
+				sb0.setSplitDag(true);
 				ret.add(sb0);
 
 				//maintain rewrite status
@@ -123,4 +123,10 @@ public ArrayList<StatementBlock> rewriteStatementBlock(StatementBlock sb, Progra
 
 		return ret;
 	}
+	
+	@Override
+	public List<StatementBlock> rewriteStatementBlocks(List<StatementBlock> sbs, 
+			ProgramRewriteStatus sate) throws HopsException {
+		return sbs;
+	}
 }