From a7fadd185ae603bcaa093b56b7812a83d5484420 Mon Sep 17 00:00:00 2001 From: ywcb00 Date: Mon, 26 Oct 2026 17:05:48 +0100 Subject: [PATCH 1/4] feat(main/hops/rewrite/SparsityOptimizer.java): create separate rewrite class sparsity optimizer to perform complete rewrites of isolated mm chain statement blocks --- .../sysds/hops/rewrite/ProgramRewriter.java | 2 + .../RewriteSplitDagSparsityRecompile.java | 23 +-- .../sysds/hops/rewrite/SparsityOptimizer.java | 158 ++++++++++++++++++ .../RewriteSplitDagSparsityRecompileTest.java | 14 +- 4 files changed, 171 insertions(+), 26 deletions(-) create mode 100644 src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java 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 178a3ecf238..6257b3563a7 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java @@ -109,6 +109,8 @@ public ProgramRewriter(boolean staticRewrites, boolean dynamicRewrites) if( OptimizerUtils.ALLOW_SPLIT_HOP_DAGS && ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_RECOMPILE) ) _sbRuleSet.add(new RewriteSplitDagSparsityRecompile()); + if( ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_RECOMPILE) ) + _sbRuleSet.add(new SparsityOptimizer()); 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/RewriteSplitDagSparsityRecompile.java b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java index f0efaaf2e4f..60598d3ef27 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/RewriteSplitDagSparsityRecompile.java @@ -20,13 +20,10 @@ 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; @@ -54,10 +51,6 @@ private static Hop getFirstTransientWriteParent(Hop hop) { } 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 @@ -158,23 +151,19 @@ private void collectNnzComputationCandidates(ArrayList roots, ArrayList HopRewriteUtils.isMatrixMultiply(parentHop)); } - public static boolean allParentsAreMatrixMultiply(Hop hop) { + private static boolean allParentsAreMatrixMultiply(Hop hop) { return hop.getParent().stream().allMatch(parentHop -> HopRewriteUtils.isMatrixMultiply(parentHop)); } - public static boolean anyChildIsMatrixMultiply(Hop hop) { + private 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) { + private static boolean isPartOfMMChain(Hop hop) { // hop is part of mm chain with at least two matrix multplications return HopRewriteUtils.isMatrixMultiply(hop) && (anyParentIsMatrixMultiply(hop) || anyChildIsMatrixMultiply(hop)); @@ -184,7 +173,7 @@ public static boolean anyParentIsMMChain(Hop hop) { List parentHops = hop.getParent(); if(parentHops == null) return false; - return parentHops.stream().anyMatch(pHop -> isMMChain(pHop)); + return parentHops.stream().anyMatch(pHop -> isPartOfMMChain(pHop)); } private void rCollectNnzComputationCandidates(Hop hop, ArrayList cand, Hop parentHop) { @@ -205,7 +194,7 @@ && anyParentIsMMChain(hop)) { // split only mm chains (i.e., at least two matrix cand.add(hop); } else if((!HopRewriteUtils.isMatrixMultiply(parentHop) || !allParentsAreMatrixMultiply(hop)) - && isMMChain(hop) // split only mm chains (i.e., at least two matrix multiplications) + && isPartOfMMChain(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); diff --git a/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java new file mode 100644 index 00000000000..22c2a80c42c --- /dev/null +++ b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java @@ -0,0 +1,158 @@ +/* + * 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.List; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.apache.hadoop.shaded.org.apache.commons.lang3.NotImplementedException; +import org.apache.sysds.common.Types.OpOpData; +import org.apache.sysds.hops.AggBinaryOp; +import org.apache.sysds.hops.DataOp; +import org.apache.sysds.hops.Hop; +import org.apache.sysds.parser.StatementBlock; +import org.apache.sysds.runtime.DMLRuntimeException; +import org.apache.sysds.utils.Explain; + +public class SparsityOptimizer extends StatementBlockRewriteRule { + protected static final Log LOG = LogFactory.getLog(SparsityOptimizer.class.getName()); + + @Override + public boolean createsSplitDag() { + return false; + } + + @Override + public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { + // optimize the matrix multiplication chain if it is isolated in its own statement block + if(isMMChain(sb)) + sb = optimizeMMChain(sb); + return Arrays.asList(sb); + } + + @Override + public List rewriteStatementBlocks(List sbs, ProgramRewriteStatus state) { + return sbs; + } + + private static boolean isMMChain(List hopList) { + boolean ret = true; + ret &= (hopList != null && hopList.size() > 0); + for(Hop hop : hopList) + ret &= isMMChain(hop); + return ret; + } + + private static boolean isMMChain(Hop hop) { + if(HopRewriteUtils.isMatrixMultiply(hop)) + return isMMChain(hop.getInput()); + else if(hop instanceof DataOp) + return (((DataOp)hop).getOp() == OpOpData.TRANSIENTREAD); + return false; + } + + private static boolean allChildsAreMatrixMultiply(Hop hop) { + return hop.getInput().stream().allMatch(childHop -> HopRewriteUtils.isMatrixMultiply(childHop)); + } + + private static boolean isMMChain(StatementBlock sb) { + boolean ret = true; + for(Hop hop : sb.getHops()) { + ret &= allChildsAreMatrixMultiply(hop) && hop instanceof DataOp; + ret &= (hop.getInput().size() > 0); + if(!ret) + return false; + ret &= isMMChain(hop.getInput()); + } + return ret; + } + + public StatementBlock optimizeMMChain(StatementBlock sb) { + for(Hop rootHop : sb.getHops()) { + List mmChain = new ArrayList<>(); + List mmOperators = new ArrayList<>(); + if(rootHop instanceof DataOp && ((DataOp)rootHop).isWrite()) + prepareMMChain(rootHop.getInput(0), mmChain, mmOperators); + else + throw new DMLRuntimeException("This case should not occur and is not considered."); + + throw new NotImplementedException("TODO: Optimize the MM Chain inside the isolated statement block"); + } + return sb; + } + + public void prepareMMChain(Hop root, List mmChain, List mmOperators) { + root.resetVisitStatus(); + // initialize mmChain with inputs of root hop + mmOperators.add(root); + for( Hop hi : root.getInput() ) + mmChain.add(hi); + + // expand each hop to find the entire matrix multiplication chain + int counter = 0; + while(counter < mmChain.size()) { + boolean expandable = true; + Hop hop = mmChain.get(counter); + + /* + * Check if mmChain[i] is expandable: + * 1) It must be MATMULT + * 2) It must not have been visited already + * (one MATMULT should get expanded only in one chain) + * 3) Its output should not be used in multiple places + * (either within chain or outside the chain) + */ + expandable &= HopRewriteUtils.isMatrixMultiply(hop) && !((AggBinaryOp)hop).hasLeftPMInput(); + expandable &= !hop.isVisited(); + expandable &= !(hop.getParent().size() > 1); + + if(!expandable) { + counter += 1; + } + else { + List tempList = mmChain.get(counter).getInput(); + // add current operator to mmOperators and its input nodes to mm chain + mmOperators.add(mmChain.get(counter)); + mmChain.set(counter, tempList.get(0)); + mmChain.add(counter + 1, tempList.get(1)); + } + } + + // print the matrix multiplication chain + if(LOG.isTraceEnabled()) { + LOG.trace("Identified MM Chain: "); + for(Hop h : mmChain) { + logTraceHop(h, 1); + } + } + } + + protected void logTraceHop(Hop hop, int level) { + if(LOG.isTraceEnabled()) { + String offset = Explain.getIdentation(level); + LOG.trace(offset + "Hop " + hop.getName() + + "(" + hop.getClass().getSimpleName() + ", " + hop.getHopID() + ")" + " " + + hop.getDim1() + "x" + hop.getDim2() + " sp" + hop.getSparsity()); + } + } +} 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 ac179eb3b9e..5cc57749c2a 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 @@ -59,10 +59,10 @@ public class RewriteSplitDagSparsityRecompileTest extends AutomatedTestBase { private static final String TEST_DIR = "functions/rewrite/"; private static final String TEST_CLASS_DIR = 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; + private static final String PACKAGE_SPARSITYOPTIMIZER = "org.apache.sysds.hops.rewrite.SparsityOptimizer"; private static Level _oldLevelRecompile; + private static Level _oldLevelSparsityOptimizer; @Parameterized.Parameter(0) public int rows; @@ -93,17 +93,17 @@ public void setUp() { 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(); Logger.getLogger(PACKAGE_RECOMPILE).setLevel(Level.TRACE); + _oldLevelSparsityOptimizer = Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).getLevel(); + Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).setLevel(Level.TRACE); } @Override public void tearDown() { super.tearDown(); - Logger.getLogger(PACKAGE_REWRITE).setLevel(_oldLevelRewrite); Logger.getLogger(PACKAGE_RECOMPILE).setLevel(_oldLevelRecompile); + Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).setLevel(_oldLevelSparsityOptimizer); } @Test @@ -156,7 +156,6 @@ private void testSplitDagSparsityRecompile(String testName, int numSBs) { try { DMLConfig dmlConfig = new DMLConfig(getCurConfigFile().getPath()); - dmlConfig.setTextValue(DMLConfig.SPARSITY_REWRITES, "true"); dmlConfig.setTextValue(DMLConfig.SPARSITY_RECOMPILE, "true"); overwriteCurrentConfig(dmlConfig); } @@ -185,13 +184,10 @@ private void testSplitDagSparsityRecompile(String testName, int numSBs) { log_original.stream().map(l -> l.getMessage().toString()).toArray(String[]::new)); 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("RewriteSplitDagSparsityRecompile")); - Assert.assertFalse(log_reference_string.contains("mmchainoptsparse")); } finally { OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = oldFlag1; From 44abfc3057524d2c64f8d24e01b3372572b759ea Mon Sep 17 00:00:00 2001 From: ywcb00 Date: Wed, 28 Oct 2026 09:39:33 +0100 Subject: [PATCH 2/4] chore(main/hops/rewrite/SparsityOptimizer.java): change indentation to use tabulators consistently --- .../sysds/hops/rewrite/SparsityOptimizer.java | 120 +++++++++--------- 1 file changed, 60 insertions(+), 60 deletions(-) diff --git a/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java index 22c2a80c42c..ed3898d56cf 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java @@ -37,23 +37,23 @@ public class SparsityOptimizer extends StatementBlockRewriteRule { protected static final Log LOG = LogFactory.getLog(SparsityOptimizer.class.getName()); - @Override - public boolean createsSplitDag() { - return false; - } - - @Override - public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { - // optimize the matrix multiplication chain if it is isolated in its own statement block - if(isMMChain(sb)) - sb = optimizeMMChain(sb); - return Arrays.asList(sb); - } - - @Override - public List rewriteStatementBlocks(List sbs, ProgramRewriteStatus state) { - return sbs; - } + @Override + public boolean createsSplitDag() { + return false; + } + + @Override + public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { + // optimize the matrix multiplication chain if it is isolated in its own statement block + if(isMMChain(sb)) + sb = optimizeMMChain(sb); + return Arrays.asList(sb); + } + + @Override + public List rewriteStatementBlocks(List sbs, ProgramRewriteStatus state) { + return sbs; + } private static boolean isMMChain(List hopList) { boolean ret = true; @@ -75,7 +75,7 @@ private static boolean allChildsAreMatrixMultiply(Hop hop) { return hop.getInput().stream().allMatch(childHop -> HopRewriteUtils.isMatrixMultiply(childHop)); } - private static boolean isMMChain(StatementBlock sb) { + private static boolean isMMChain(StatementBlock sb) { boolean ret = true; for(Hop hop : sb.getHops()) { ret &= allChildsAreMatrixMultiply(hop) && hop instanceof DataOp; @@ -87,56 +87,56 @@ private static boolean isMMChain(StatementBlock sb) { return ret; } - public StatementBlock optimizeMMChain(StatementBlock sb) { - for(Hop rootHop : sb.getHops()) { - List mmChain = new ArrayList<>(); - List mmOperators = new ArrayList<>(); - if(rootHop instanceof DataOp && ((DataOp)rootHop).isWrite()) - prepareMMChain(rootHop.getInput(0), mmChain, mmOperators); - else - throw new DMLRuntimeException("This case should not occur and is not considered."); - - throw new NotImplementedException("TODO: Optimize the MM Chain inside the isolated statement block"); - } - return sb; - } - - public void prepareMMChain(Hop root, List mmChain, List mmOperators) { - root.resetVisitStatus(); - // initialize mmChain with inputs of root hop + public StatementBlock optimizeMMChain(StatementBlock sb) { + for(Hop rootHop : sb.getHops()) { + List mmChain = new ArrayList<>(); + List mmOperators = new ArrayList<>(); + if(rootHop instanceof DataOp && ((DataOp)rootHop).isWrite()) + prepareMMChain(rootHop.getInput(0), mmChain, mmOperators); + else + throw new DMLRuntimeException("This case should not occur and is not considered."); + + throw new NotImplementedException("TODO: Optimize the MM Chain inside the isolated statement block"); + } + return sb; + } + + public void prepareMMChain(Hop root, List mmChain, List mmOperators) { + root.resetVisitStatus(); + // initialize mmChain with inputs of root hop mmOperators.add(root); for( Hop hi : root.getInput() ) mmChain.add(hi); - // expand each hop to find the entire matrix multiplication chain - int counter = 0; - while(counter < mmChain.size()) { - boolean expandable = true; - Hop hop = mmChain.get(counter); + // expand each hop to find the entire matrix multiplication chain + int counter = 0; + while(counter < mmChain.size()) { + boolean expandable = true; + Hop hop = mmChain.get(counter); /* * Check if mmChain[i] is expandable: * 1) It must be MATMULT * 2) It must not have been visited already - * (one MATMULT should get expanded only in one chain) + * (one MATMULT should get expanded only in one chain) * 3) Its output should not be used in multiple places - * (either within chain or outside the chain) + * (either within chain or outside the chain) */ - expandable &= HopRewriteUtils.isMatrixMultiply(hop) && !((AggBinaryOp)hop).hasLeftPMInput(); - expandable &= !hop.isVisited(); - expandable &= !(hop.getParent().size() > 1); - - if(!expandable) { - counter += 1; - } - else { - List tempList = mmChain.get(counter).getInput(); - // add current operator to mmOperators and its input nodes to mm chain - mmOperators.add(mmChain.get(counter)); - mmChain.set(counter, tempList.get(0)); - mmChain.add(counter + 1, tempList.get(1)); - } - } + expandable &= HopRewriteUtils.isMatrixMultiply(hop) && !((AggBinaryOp)hop).hasLeftPMInput(); + expandable &= !hop.isVisited(); + expandable &= !(hop.getParent().size() > 1); + + if(!expandable) { + counter += 1; + } + else { + List tempList = mmChain.get(counter).getInput(); + // add current operator to mmOperators and its input nodes to mm chain + mmOperators.add(mmChain.get(counter)); + mmChain.set(counter, tempList.get(0)); + mmChain.add(counter + 1, tempList.get(1)); + } + } // print the matrix multiplication chain if(LOG.isTraceEnabled()) { @@ -145,14 +145,14 @@ public void prepareMMChain(Hop root, List mmChain, List mmOperators) { logTraceHop(h, 1); } } - } + } protected void logTraceHop(Hop hop, int level) { if(LOG.isTraceEnabled()) { String offset = Explain.getIdentation(level); LOG.trace(offset + "Hop " + hop.getName() + - "(" + hop.getClass().getSimpleName() + ", " + hop.getHopID() + ")" + " " + - hop.getDim1() + "x" + hop.getDim2() + " sp" + hop.getSparsity()); + "(" + hop.getClass().getSimpleName() + ", " + hop.getHopID() + ")" + " " + + hop.getDim1() + "x" + hop.getDim2() + " sp" + hop.getSparsity()); } } } From 5b8090c18f74f1cc9a0e23914c5ea5d57243ed95 Mon Sep 17 00:00:00 2001 From: ywcb00 Date: Tue, 3 Nov 2026 12:05:46 +0100 Subject: [PATCH 3/4] fix(main/hops/rewrite/SparsityOptimizer.java): change sparsity optimizer to being a hop rewrite rule instead of statement block rewrite NOTE: hop rewrite rules are also executed after computing the previous statementblock feat(main/hops/rewrite/SparsityOptimizer.java): replicate the optimization from normal sparsity-based mm chain rewrite create the method to "sculpt" a mm chain, which is allowed to completely rewrite the mm chain and ingest a new hop structure feat(test/functions/rewrite/SparsityOptimizerTest.java): create preliminary junit test for the sparsity optimizer --- .../sysds/hops/rewrite/ProgramRewriter.java | 5 +- .../sysds/hops/rewrite/SparsityOptimizer.java | 88 ++++---- .../rewrite/SparsityOptimizerTest.java | 191 ++++++++++++++++++ .../rewrite/SparsityOptimizerMMChain.dml | 30 +++ 4 files changed, 264 insertions(+), 50 deletions(-) create mode 100644 src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java create mode 100644 src/test/scripts/functions/rewrite/SparsityOptimizerMMChain.dml 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 6257b3563a7..849f2d91e80 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/ProgramRewriter.java @@ -109,8 +109,6 @@ public ProgramRewriter(boolean staticRewrites, boolean dynamicRewrites) if( OptimizerUtils.ALLOW_SPLIT_HOP_DAGS && ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_RECOMPILE) ) _sbRuleSet.add(new RewriteSplitDagSparsityRecompile()); - if( ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_RECOMPILE) ) - _sbRuleSet.add(new SparsityOptimizer()); if( OptimizerUtils.ALLOW_SPLIT_HOP_DAGS && ConfigurationManager.getCompilerConfigFlag(ConfigType.ALLOW_INDIVIDUAL_SB_SPECIFIC_OPS) ) _sbRuleSet.add( new RewriteSplitDagDataDependentOperators() ); //dependency: merge blocks @@ -151,6 +149,9 @@ public ProgramRewriter(boolean staticRewrites, boolean dynamicRewrites) if(ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_REWRITES)) { _dagRuleSet.add( new RewriteMatrixMultChainOptimizationSparse() ); } + if( ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_RECOMPILE) ) { + _dagRuleSet.add(new SparsityOptimizer()); + } if( OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION ) { _dagRuleSet.add( new RewriteAlgebraicSimplificationDynamic() ); //dependencies: cse _dagRuleSet.add( new RewriteAlgebraicSimplificationStatic() ); //dependencies: cse diff --git a/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java index ed3898d56cf..71c3236f40d 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java @@ -20,39 +20,36 @@ package org.apache.sysds.hops.rewrite; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; -import org.apache.hadoop.shaded.org.apache.commons.lang3.NotImplementedException; import org.apache.sysds.common.Types.OpOpData; import org.apache.sysds.hops.AggBinaryOp; import org.apache.sysds.hops.DataOp; import org.apache.sysds.hops.Hop; -import org.apache.sysds.parser.StatementBlock; -import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.utils.Explain; -public class SparsityOptimizer extends StatementBlockRewriteRule { +public class SparsityOptimizer extends HopRewriteRule { protected static final Log LOG = LogFactory.getLog(SparsityOptimizer.class.getName()); @Override - public boolean createsSplitDag() { - return false; - } - - @Override - public List rewriteStatementBlock(StatementBlock sb, ProgramRewriteStatus state) { - // optimize the matrix multiplication chain if it is isolated in its own statement block - if(isMMChain(sb)) - sb = optimizeMMChain(sb); - return Arrays.asList(sb); + public ArrayList rewriteHopDAGs(ArrayList roots, ProgramRewriteStatus state) { + ArrayList ret = new ArrayList(); + for(Hop rootHop : roots) + ret.add(rewriteHopDAG(rootHop, state)); + return ret; } @Override - public List rewriteStatementBlocks(List sbs, ProgramRewriteStatus state) { - return sbs; + public Hop rewriteHopDAG(Hop root, ProgramRewriteStatus state) { + if(root.isVisited()) + return root; + + if(root instanceof DataOp && ((DataOp)root).isWrite() && root.getInput().size() > 0 && + HopRewriteUtils.isMatrixMultiply(root.getInput(0)) && isMMChain(root.getInput())) + return optimizeMMChain(root); + return root; } private static boolean isMMChain(List hopList) { @@ -71,34 +68,27 @@ else if(hop instanceof DataOp) return false; } - private static boolean allChildsAreMatrixMultiply(Hop hop) { - return hop.getInput().stream().allMatch(childHop -> HopRewriteUtils.isMatrixMultiply(childHop)); - } + public Hop optimizeMMChain(Hop rootHop) { + List mmChain = new ArrayList<>(); + List mmOperators = new ArrayList<>(); + Hop mmHop = rootHop.getInput(0); + prepareMMChain(mmHop, mmChain, mmOperators); - private static boolean isMMChain(StatementBlock sb) { - boolean ret = true; - for(Hop hop : sb.getHops()) { - ret &= allChildsAreMatrixMultiply(hop) && hop instanceof DataOp; - ret &= (hop.getInput().size() > 0); - if(!ret) - return false; - ret &= isMMChain(hop.getInput()); + Hop newMMHop = sculptMMChain(mmHop, mmChain, mmOperators); + if(newMMHop != null) { + HopRewriteUtils.replaceChildReference(rootHop, mmHop, newMMHop); + mmHop = newMMHop; } - return ret; - } - public StatementBlock optimizeMMChain(StatementBlock sb) { - for(Hop rootHop : sb.getHops()) { - List mmChain = new ArrayList<>(); - List mmOperators = new ArrayList<>(); - if(rootHop instanceof DataOp && ((DataOp)rootHop).isWrite()) - prepareMMChain(rootHop.getInput(0), mmChain, mmOperators); - else - throw new DMLRuntimeException("This case should not occur and is not considered."); - - throw new NotImplementedException("TODO: Optimize the MM Chain inside the isolated statement block"); + // print the matrix multiplication chain + if(LOG.isTraceEnabled()) { + LOG.trace("SparsityOptimizer: Identified and optimized MM Chain: "); + for(Hop h : mmChain) { + logTraceHop(h, 1); + } } - return sb; + + return rootHop; } public void prepareMMChain(Hop root, List mmChain, List mmOperators) { @@ -126,6 +116,8 @@ public void prepareMMChain(Hop root, List mmChain, List mmOperators) { expandable &= !hop.isVisited(); expandable &= !(hop.getParent().size() > 1); + hop.setVisited(); + if(!expandable) { counter += 1; } @@ -137,14 +129,14 @@ public void prepareMMChain(Hop root, List mmChain, List mmOperators) { mmChain.add(counter + 1, tempList.get(1)); } } + } - // print the matrix multiplication chain - if(LOG.isTraceEnabled()) { - LOG.trace("Identified MM Chain: "); - for(Hop h : mmChain) { - logTraceHop(h, 1); - } - } + // Optimize the matrix multiplication chain by rewriting it + // This method can modify the mm chain completely, even re-creating it from scratch is allowed. + // NOTE: This method should return the new mm hop when re-creating it from scratch. + private Hop sculptMMChain(Hop mmHop, List mmChain, List mmOperators) { + (new RewriteMatrixMultChainOptimizationSparse()).optimizeMMChain(mmHop, mmChain, mmOperators, null); + return null; } protected void logTraceHop(Hop hop, int level) { diff --git a/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java b/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java new file mode 100644 index 00000000000..0a4c273162e --- /dev/null +++ b/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java @@ -0,0 +1,191 @@ +/* + * 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.test.functions.rewrite; + +import org.apache.log4j.Level; +import org.apache.log4j.Logger; +import org.apache.log4j.spi.LoggingEvent; +import org.apache.sysds.conf.ConfigurationManager; +import org.apache.sysds.conf.DMLConfig; +import org.apache.sysds.hops.OptimizerUtils; +import org.apache.sysds.hops.recompile.Recompiler; +import org.apache.sysds.runtime.matrix.data.MatrixValue; +import org.apache.sysds.test.AutomatedTestBase; +import org.apache.sysds.test.TestConfiguration; +import org.apache.sysds.test.TestUtils; +import org.apache.sysds.test.LoggingUtils; +import org.apache.sysds.test.LoggingUtils.TestAppender; + +import org.junit.Assert; +import org.junit.runners.Parameterized; +import org.junit.Test; +import org.junit.runner.RunWith; + +import java.io.FileNotFoundException; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.stream.DoubleStream; +import java.util.stream.Stream; + +@RunWith(value = Parameterized.class) +@net.jcip.annotations.NotThreadSafe +public class SparsityOptimizerTest extends AutomatedTestBase { + + private static final String TEST_NAME_MMCHAIN = "SparsityOptimizerMMChain"; + private static final String TEST_DIR = "functions/rewrite/"; + private static final String TEST_CLASS_DIR = + TEST_DIR + SparsityOptimizerTest.class.getSimpleName() + "/"; + private static final String PACKAGE_RECOMPILE = "org.apache.sysds.hops.rewrite.StatementBlockRewriteRule"; + private static Level _oldLevelRecompile; + private static final String PACKAGE_SPARSITYOPTIMIZER = "org.apache.sysds.hops.rewrite.SparsityOptimizer"; + private static Level _oldLevelSparsityOptimizer; + + @Parameterized.Parameter(0) + public int rows; + + @Parameterized.Parameter(1) + public int cols; + + @Parameterized.Parameter(2) + public double[] sparsities; + + @Parameterized.Parameter(3) + public double eps; + + @Parameterized.Parameter(4) + public boolean tsmm; + + @Parameterized.Parameters + public static Collection data() { + return Arrays.asList(new Object[][] { + // {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},}); + } + + @Override + public void setUp() { + TestUtils.clearAssertionInformation(); + addTestConfiguration(TEST_NAME_MMCHAIN, new TestConfiguration(TEST_CLASS_DIR, TEST_NAME_MMCHAIN, new String[] {"R"})); + _oldLevelRecompile = Logger.getLogger(PACKAGE_RECOMPILE).getLevel(); + Logger.getLogger(PACKAGE_RECOMPILE).setLevel(Level.TRACE); + _oldLevelSparsityOptimizer = Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).getLevel(); + Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).setLevel(Level.TRACE); + } + + @Override + public void tearDown() { + super.tearDown(); + Logger.getLogger(PACKAGE_RECOMPILE).setLevel(_oldLevelRecompile); + Logger.getLogger(PACKAGE_SPARSITYOPTIMIZER).setLevel(_oldLevelSparsityOptimizer); + } + + @Test + public void testSparsityOptimizationMMChain() { + testSparsityOptimization(TEST_NAME_MMCHAIN); + } + + private void testSparsityOptimization(String testName) { + boolean oldFlag = OptimizerUtils.ALLOW_SPLIT_HOP_DAGS; + DMLConfig oldDMLConfig = ConfigurationManager.getDMLConfig(); + + try { + TestConfiguration config = getTestConfiguration(testName); + loadTestConfiguration(config); + + String HOME = SCRIPT_DIR + TEST_DIR; + fullDMLScriptName = HOME + testName + ".dml"; + programArgs = new String[] {"-explain", "hops", "-stats", + "-args", input("X"), input("Y"), output("R")}; + + 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); + long Y_nnz = Stream.of(Y).mapToLong(row -> DoubleStream.of(row).filter(val -> val != 0).count()).sum(); + writeInputMatrixWithMTD("X", X, true); // write matrix characteristics w/ missing nnz + writeInputMatrixWithMTD("Y", Y, Y_nnz, true); + + // run reference test + TestAppender appender = LoggingUtils.overwrite(); // capture log output + runTest(true, false, null, -1); + List log_reference = LoggingUtils.reinsert(appender); // revert the logger to print to stdout + + // get reference result matrix + HashMap referenceResult = readDMLMatrixFromOutputDir("R"); + + try { + DMLConfig dmlConfig = new DMLConfig(getCurConfigFile().getPath()); + dmlConfig.setTextValue(DMLConfig.SPARSITY_RECOMPILE, "true"); + overwriteCurrentConfig(dmlConfig); + } + catch(FileNotFoundException fnfe) { + Assert.fail( + "Could not find DML config file: " + getCurConfigFile().getPath() + " . " + fnfe.getMessage()); + } + catch(IOException ioe) { + Assert.fail("Could not overwrite the DML configuration file. " + ioe.getMessage()); + } + + // run original test + appender = LoggingUtils.overwrite(); // capture log output + runTest(true, false, null, -1); + List log_original = LoggingUtils.reinsert(appender); // revert the logger to print to stdout + + // get resulting matrix + HashMap originalResult = readDMLMatrixFromOutputDir("R"); + + // compare matrices + TestUtils.compareMatrices(originalResult, referenceResult, eps, "Stat-Original", "Stat-Reference"); + + final String DELIMITER = "\n"; + // check original log output + String log_original_string = String.join(DELIMITER, + log_original.stream().map(l -> l.getMessage().toString()).toArray(String[]::new)); + System.out.println(log_original_string); + Assert.assertTrue(log_original_string.contains("RewriteSplitDagSparsityRecompile: Split DAG into")); + Assert.assertTrue(log_original_string.contains("SparsityOptimizer: Identified and optimized 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("RewriteSplitDagSparsityRecompile")); + Assert.assertFalse(log_reference_string.contains("SparsityOptimizer")); + } + finally { + OptimizerUtils.ALLOW_SPLIT_HOP_DAGS = oldFlag; + try { + overwriteCurrentConfig(oldDMLConfig); + } + catch(IOException ioe) { + Assert.fail("Unable to restore the previous DML configuration. " + ioe.getMessage()); + } + Recompiler.reinitRecompiler(); + } + } + + private void overwriteCurrentConfig(DMLConfig config) throws IOException { + Files.write(getCurConfigFile().toPath(), config.serializeDMLConfig().getBytes(StandardCharsets.UTF_8)); + } +} diff --git a/src/test/scripts/functions/rewrite/SparsityOptimizerMMChain.dml b/src/test/scripts/functions/rewrite/SparsityOptimizerMMChain.dml new file mode 100644 index 00000000000..2b4308b51f7 --- /dev/null +++ b/src/test/scripts/functions/rewrite/SparsityOptimizerMMChain.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); + +S = t(X) %*% X %*% Y; + +R = t(S); + +write(R, $3); From c500b5da381fb7ffb9fb900e5bc097e0dc767286 Mon Sep 17 00:00:00 2001 From: ywcb00 Date: Tue, 3 Nov 2026 15:26:47 +0100 Subject: [PATCH 4/4] chore(main/hops/rewrite/SparsityOptimizer.java): log the sparsity similar to the mm chain rewrite class chore(test/functions/rewrite/SparsityOptimizerTest.java): set the estimator type explicitly --- .../java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java | 2 +- .../sysds/test/functions/rewrite/SparsityOptimizerTest.java | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java index 71c3236f40d..36b84b9cf9f 100644 --- a/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java +++ b/src/main/java/org/apache/sysds/hops/rewrite/SparsityOptimizer.java @@ -144,7 +144,7 @@ protected void logTraceHop(Hop hop, int level) { String offset = Explain.getIdentation(level); LOG.trace(offset + "Hop " + hop.getName() + "(" + hop.getClass().getSimpleName() + ", " + hop.getHopID() + ")" + " " + - hop.getDim1() + "x" + hop.getDim2() + " sp" + hop.getSparsity()); + hop.getDim1() + "x" + hop.getDim2() + " sparsity" + hop.getSparsity()); } } } diff --git a/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java b/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java index 0a4c273162e..48c7046da3e 100644 --- a/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java +++ b/src/test/java/org/apache/sysds/test/functions/rewrite/SparsityOptimizerTest.java @@ -139,6 +139,7 @@ private void testSparsityOptimization(String testName) { try { DMLConfig dmlConfig = new DMLConfig(getCurConfigFile().getPath()); dmlConfig.setTextValue(DMLConfig.SPARSITY_RECOMPILE, "true"); + dmlConfig.setTextValue(DMLConfig.SPARSITY_ESTIMATOR, "BASIC_AVG"); overwriteCurrentConfig(dmlConfig); } catch(FileNotFoundException fnfe) {