From da4cb624d4245caf955f19ca6024984cdb741a4c Mon Sep 17 00:00:00 2001 From: ywcb00 Date: Thu, 29 Oct 2026 17:08:31 +0100 Subject: [PATCH] fix(main/hops/rewrite/RewriteSplitDagSparsityRecompile.java): reuse existing transient writes instead of introducing redundant transient reads and writes --- .../RewriteSplitDagSparsityRecompile.java | 31 +++++++++++++++---- .../RewriteSplitDagSparsityRecompileTest.java | 7 +++++ ...writeSplitDagSparsityRecompileMMChain2.dml | 30 ++++++++++++++++++ 3 files changed, 62 insertions(+), 6 deletions(-) create mode 100644 src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileMMChain2.dml 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 c3237677ce9..f0efaaf2e4f 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java @@ -45,6 +45,14 @@ public List 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 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 rRewriteStatementBlock(StatementBlock sb, ProgramRew ArrayList 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 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) 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 1802ef54660..ac179eb3b9e 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 @@ -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; diff --git a/src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileMMChain2.dml b/src/test/scripts/functions/rewrite/RewriteSplitDagSparsityRecompileMMChain2.dml new file mode 100644 index 00000000000..f178334b62f --- /dev/null +++ b/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);