[SYSTEMDS-3979] Reuse Existing Transient Writes when Splitting DAG

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

@@ -45,6 +45,14 @@ public List<StatementBlock> rewriteStatementBlock(StatementBlock sb, ProgramRewr
 		return rRewriteStatementBlock(sb, state);
 	}
 
+	private static Hop getFirstTransientWriteParent(Hop hop) {
+		for(Hop p : hop.getParent()) {
+			if(p instanceof DataOp && ((DataOp)p).getOp() == OpOpData.TRANSIENTWRITE)
+				return p;
+		}
+		return null;
+	}
+
 	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))
@@ -78,21 +86,32 @@ public List<StatementBlock> rRewriteStatementBlock(StatementBlock sb, ProgramRew
 			ArrayList<Hop> sbSplitHops = new ArrayList<>();
 			for(Hop cHop : cand) {
 				// create new transient read
-				String varname = createCutVarName(false);
+				String varname;
+				DataOp tWrite = (DataOp)getFirstTransientWriteParent(cHop);
+				if(cHop instanceof DataOp && ((DataOp)cHop).isRead() && tWrite != null) {
+					varname = tWrite.getName();
+				}
+				else {
+					varname = createCutVarName(false);
+				}
 				DataOp tRead = HopRewriteUtils.createTransientRead(varname, cHop);
 				tRead.setRequiresRecompile();
 
-				// replace persistent read with transient read
+				// replace hop with transient read operation
 				ArrayList<Hop> 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);
+					if(!candChilds.contains(parent)) { // anomaly filter
+						if(parent != tWrite)
+							HopRewriteUtils.replaceChildReference(parent, cHop, tRead);
+						else
+							sb.getHops().remove(parent);
+					}
 				}
 
-				// add reblock sub dag to first statement block
-				DataOp tWrite = HopRewriteUtils.createTransientWrite(varname, cHop);
+				if(tWrite == null)
+					tWrite = HopRewriteUtils.createTransientWrite(varname, cHop);
 				sbSplitHops.add(tWrite);
 
 				// update live in and out of new statement block (for piggybacking)

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

@@ -55,6 +55,7 @@ public class RewriteSplitDagSparsityRecompileTest extends AutomatedTestBase {
 
 	private static final String TEST_NAME_READ = "RewriteSplitDagSparsityRecompileRead";
 	private static final String TEST_NAME_MMCHAIN = "RewriteSplitDagSparsityRecompileMMChain";
+	private static final String TEST_NAME_MMCHAIN2 = "RewriteSplitDagSparsityRecompileMMChain2";
 	private static final String TEST_DIR = "functions/rewrite/";
 	private static final String TEST_CLASS_DIR =
 		TEST_DIR + RewriteSplitDagSparsityRecompileTest.class.getSimpleName() + "/";
@@ -91,6 +92,7 @@ public void setUp() {
 		TestUtils.clearAssertionInformation();
 		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();
@@ -114,6 +116,11 @@ public void testSplitDagSparsityRecompileMMChain() {
 		testSplitDagSparsityRecompile(TEST_NAME_MMCHAIN, 4);
 	}
 
+	@Test
+	public void testSplitDagSparsityRecompileMMChain2() {
+		testSplitDagSparsityRecompile(TEST_NAME_MMCHAIN2, 4);
+	}
+
 	private void testSplitDagSparsityRecompile(String testName, int numSBs) {
 		boolean oldFlag1 = OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES;
 		boolean oldFlag2 = OptimizerUtils.ALLOW_SUM_PRODUCT_REWRITES;

src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileMMChain2.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);
+
+X0 = (X == 0);
+
+R = X %*% t(X) %*% X0 %*% t(X) %*% X %*% t(X)
+
+write(R, $3);