[SYSTEMDS-3958] Add XML options for sparse rewrites

src/main/java/org/apache/sysds/conf/DMLConfig.java

@@ -43,6 +43,7 @@
 import org.apache.sysds.hops.codegen.SpoofCompiler.CompilerType;
 import org.apache.sysds.hops.codegen.SpoofCompiler.GeneratorAPI;
 import org.apache.sysds.hops.codegen.SpoofCompiler.PlanSelector;
+import org.apache.sysds.hops.estim.EstimationUtils.EstimatorType;
 import org.apache.sysds.hops.fedplanner.FTypes.FederatedPlanner;
 import org.apache.sysds.lops.Compression;
 import org.apache.sysds.lops.compile.linearization.IDagLinearizerFactory.DagLinearizer;
@@ -97,6 +98,8 @@ public class DMLConfig
 	public static final String NATIVE_BLAS          = "sysds.native.blas";
 	public static final String NATIVE_BLAS_DIR      = "sysds.native.blas.directory";
 	public static final String DAG_LINEARIZATION    = "sysds.compile.linearization";
+	public static final String SPARSITY_REWRITES    = "sysds.rewrites.sparsity.enabled"; // boolean
+	public static final String SPARSITY_ESTIMATOR   = "sysds.rewrites.sparsity.estimator"; // see EstiamtionUtils.EstimatorType
 	public static final String CODEGEN              = "sysds.codegen.enabled"; //boolean
 	public static final String CODEGEN_API          = "sysds.codegen.api"; // see SpoofCompiler.API
 	public static final String CODEGEN_COMPILER     = "sysds.codegen.compiler"; //see SpoofCompiler.CompilerType
@@ -188,6 +191,8 @@ public class DMLConfig
 		_defaultVals.put(COMPRESSED_TRANSPOSE,   "auto");
 		_defaultVals.put(COMPRESSED_TRANSFORMENCODE, "false");
 		_defaultVals.put(DAG_LINEARIZATION,      DagLinearizer.DEPTH_FIRST.name());
+		_defaultVals.put(SPARSITY_REWRITES,      "false");
+		_defaultVals.put(SPARSITY_ESTIMATOR,     EstimatorType.BASIC_AVG.name());
 		_defaultVals.put(CODEGEN,                "false" );
 		_defaultVals.put(CODEGEN_API,            GeneratorAPI.JAVA.name() );
 		_defaultVals.put(CODEGEN_COMPILER,       CompilerType.AUTO.name() );
@@ -476,6 +481,7 @@ public String getConfigInfo()  {
 			COMPRESSED_LINALG, COMPRESSED_LOSSY, COMPRESSED_VALID_COMPRESSIONS, COMPRESSED_OVERLAPPING,
 			COMPRESSED_SAMPLING_RATIO, COMPRESSED_SOFT_REFERENCE_COUNT,
 			COMPRESSED_COCODE, COMPRESSED_TRANSPOSE, COMPRESSED_TRANSFORMENCODE, DAG_LINEARIZATION,
+			SPARSITY_REWRITES, SPARSITY_ESTIMATOR,
 			CODEGEN, CODEGEN_API, CODEGEN_COMPILER, CODEGEN_OPTIMIZER, CODEGEN_PLANCACHE, CODEGEN_LITERALS,
 			STATS_MAX_WRAP_LEN, LINEAGECACHESPILL, COMPILERASSISTED_RW, BUFFERPOOL_LIMIT, MEMORY_MANAGER,
 			PRINT_GPU_MEMORY_INFO, AVAILABLE_GPUS, SYNCHRONIZE_GPU, EAGER_CUDA_FREE, GPU_RULE_BASED_PLACEMENT,

src/main/java/org/apache/sysds/hops/OptimizerUtils.java

@@ -200,10 +200,9 @@ public enum MemoryManager {
 	public static boolean ALLOW_SUM_PRODUCT_REWRITES2 = true;
 
 	/**
-	 * Enables additional mmchain optimizations. In the future, this might be merged with
-	 * ALLOW_SUM_PRODUCT_REWRITES.
+	 * Enables transpose mmchain optimizations. In the future, this might be merged with ALLOW_SUM_PRODUCT_REWRITES.
 	 */
-	public static boolean ALLOW_ADVANCED_MMCHAIN_REWRITES = false;
+	public static boolean ALLOW_TRANSPOSE_MMCHAIN_REWRITES = false;
 
 	/**
 	 * Enables a DPSize inspired algorithm rewrite for MMChain with transposes

src/main/java/org/apache/sysds/hops/estim/EstimationUtils.java

@@ -30,6 +30,55 @@
 
 public abstract class EstimationUtils 
 {
+	/**
+	 * Enumeration for the sparsity estimators supported
+	 */
+	public enum EstimatorType {
+		BASIC_AVG,
+		BASIC_WORST,
+		BITSET_MM,
+		DM,
+		LG,
+		MNC,
+		MNC_LIM,
+		MNC_EXT,
+		RS,
+		SAMPLE,
+		SAMPLE_RA;
+
+		/**
+		 * @return a sparsity estimator object corresponding to this estimator type
+		 */
+		public SparsityEstimator getEstimator() {
+			switch(this) {
+				case BASIC_AVG:
+					return new EstimatorBasicAvg();
+				case BASIC_WORST:
+					return new EstimatorBasicWorst();
+				case BITSET_MM:
+					return new EstimatorBitsetMM();
+				case DM:
+					return new EstimatorDensityMap();
+				case LG:
+					return new EstimatorLayeredGraph();
+				case MNC:
+					return new EstimatorMatrixHistogram();
+				case MNC_LIM:
+					return new EstimatorMatrixHistogram(false);
+				case MNC_EXT:
+					return new EstimatorMatrixHistogram(true);
+				case RS:
+					return new EstimatorRowWise();
+				case SAMPLE:
+					return new EstimatorSample();
+				case SAMPLE_RA:
+					return new EstimatorSampleRa();
+				default:
+					throw new DMLRuntimeException("Unknown sparsity estimator " + this.toString());
+			}
+		}
+	}
+
 	/**
 	 * This utility function computes the exact output nnz
 	 * of a self matrix product without need to materialize

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

@@ -24,6 +24,7 @@
 
 import org.apache.sysds.api.DMLScript;
 import org.apache.sysds.conf.ConfigurationManager;
+import org.apache.sysds.conf.DMLConfig;
 import org.apache.sysds.conf.CompilerConfig.ConfigType;
 import org.apache.sysds.hops.Hop;
 import org.apache.sysds.hops.OptimizerUtils;
@@ -139,9 +140,11 @@ public ProgramRewriter(boolean staticRewrites, boolean dynamicRewrites)
 			if( OptimizerUtils.ALLOW_NEW_MMCHAIN_REWRITE ) {
 				_dagRuleSet.add( new RewriteMatrixMultChainWithTransOptimization()		);
 			}
-			if(OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES){
+			if(OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES){
 				_dagRuleSet.add( new RewriteMatrixMultChainOptimizationTranspose()      ); //dependency: cse
-				_dagRuleSet.add( new RewriteMatrixMultChainOptimizationSparse()         ); //dependency: cse
+			}
+			if(ConfigurationManager.getDMLConfig().getBooleanValue(DMLConfig.SPARSITY_REWRITES)) {
+				_dagRuleSet.add( new RewriteMatrixMultChainOptimizationSparse()         );
 			}
 			if( OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION ) {
 				_dagRuleSet.add( new RewriteAlgebraicSimplificationDynamic()      ); //dependencies: cse

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

@@ -23,10 +23,13 @@
 import java.util.List;
 
 import org.apache.commons.lang3.mutable.MutableInt;
+import org.apache.sysds.conf.ConfigurationManager;
+import org.apache.sysds.conf.DMLConfig;
 import org.apache.sysds.hops.Hop;
 import org.apache.sysds.hops.OptimizerUtils;
 import org.apache.sysds.hops.estim.MMNode;
-import org.apache.sysds.hops.estim.EstimatorBasicAvg;
+import org.apache.sysds.hops.estim.SparsityEstimator;
+import org.apache.sysds.hops.estim.EstimationUtils.EstimatorType;
 import org.apache.sysds.hops.estim.SparsityEstimator.OpCode;
 
 /**
@@ -35,7 +38,7 @@
  * 
  * Solution: Classic Dynamic Programming
  * Approach: Currently, the approach based only on matrix dimensions
- * and sparsity estimates using the MNC sketch
+ * and sparsity estimates using the basic average estimator
  * Goal: To reduce the number of computations in the run-time
  * (map-reduce) layer
  */
@@ -85,9 +88,10 @@ private static int[][] mmChainDPSparse(double[] dimArray, MMNode[] sketchArray,
 		}
 
 		//compute cost-optimal chains for increasing chain sizes 
-		EstimatorBasicAvg estim = new EstimatorBasicAvg();
-		for( int l = 2; l <= size; l++ ) { // chain length
-			for( int i = 0; i < size - l + 1; i++ ) {
+		SparsityEstimator estim = EstimatorType.valueOf(ConfigurationManager.getDMLConfig()
+			.getTextValue(DMLConfig.SPARSITY_ESTIMATOR)).getEstimator();
+		for(int l = 2; l <= size; l++) { // chain length
+			for(int i = 0; i < size - l + 1; i++) {
 				int j = i + l - 1;
 				// find cost of (i,j)
 				dpMatrix[i][j] = Double.MAX_VALUE;

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

@@ -21,7 +21,15 @@
 
 import org.junit.Assert;
 import org.junit.Test;
+
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+
 import org.apache.sysds.common.Types.ExecMode;
+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.test.AutomatedTestBase;
@@ -123,19 +131,33 @@ public void setUp() {
 
 	private void runTestMatrixChainDP(String testName) {
 		ExecMode platformOld = rtplatform;
-		boolean rewritesOld = OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION;
-		boolean newMMchain1 = OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES;
-		boolean newMMchain2 = OptimizerUtils.ALLOW_NEW_MMCHAIN_REWRITE;
+		boolean oldFlag1 = OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION;
+		boolean oldFlag2 = OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES;
+		boolean oldFlag3 = OptimizerUtils.ALLOW_NEW_MMCHAIN_REWRITE;
+		DMLConfig oldDMLConfig = ConfigurationManager.getDMLConfig();
 
 		try {
 			rtplatform = ExecMode.SINGLE_NODE;
 			OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION = true;
-			OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES = true;
+			OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = true;
 			OptimizerUtils.ALLOW_NEW_MMCHAIN_REWRITE = true;
 
 			TestConfiguration config = getTestConfiguration(testName);
 			loadTestConfiguration(config);
 
+			try {
+				DMLConfig dmlConfig = new DMLConfig(getCurConfigFile().getPath());
+				dmlConfig.setTextValue(DMLConfig.SPARSITY_REWRITES, "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());
+			}
+
 			String HOME = SCRIPT_DIR + TEST_DIR;
 			fullDMLScriptName = HOME + testName + ".dml";
 
@@ -300,11 +322,21 @@ private void runTestMatrixChainDP(String testName) {
 				}
 			}
 		} finally {
-			OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION = rewritesOld;
-			OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES = newMMchain1;
-			OptimizerUtils.ALLOW_NEW_MMCHAIN_REWRITE = newMMchain2;
+			OptimizerUtils.ALLOW_ALGEBRAIC_SIMPLIFICATION = oldFlag1;
+			OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = oldFlag2;
+			OptimizerUtils.ALLOW_NEW_MMCHAIN_REWRITE = oldFlag3;
+			try {
+				overwriteCurrentConfig(oldDMLConfig);
+			}
+			catch(IOException ioe) {
+				Assert.fail("Unable to restore the previous DML configuration. " + ioe.getMessage());
+			}
 			rtplatform = platformOld;
 			Recompiler.reinitRecompiler();
 		}
 	}
+
+	private void overwriteCurrentConfig(DMLConfig config) throws IOException {
+		Files.write(getCurConfigFile().toPath(), config.serializeDMLConfig().getBytes(StandardCharsets.UTF_8));
+	}
 }

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

@@ -23,6 +23,8 @@
 import org.apache.log4j.Logger;
 import org.apache.log4j.spi.LoggingEvent;
 import org.apache.sysds.common.Opcodes;
+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;
@@ -37,6 +39,10 @@
 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;
@@ -103,21 +109,37 @@ public void testMatrixMultChainOptSparseRewrites() {
 	}
 
 	private void testRewriteMatrixMultChainOpSparse(boolean rewrites) {
-		boolean oldFlag1 = OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES;
+		boolean oldFlag1 = OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES;
 		boolean oldFlag2 = OptimizerUtils.ALLOW_SUM_PRODUCT_REWRITES;
+		DMLConfig oldDMLConfig = ConfigurationManager.getDMLConfig();
 
 		try {
 			TestConfiguration config = getTestConfiguration(TEST_NAME);
 			loadTestConfiguration(config);
 
+			try {
+				DMLConfig dmlConfig = new DMLConfig(getCurConfigFile().getPath());
+				dmlConfig.setTextValue(DMLConfig.SPARSITY_REWRITES, String.valueOf(rewrites));
+				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());
+			}
+
 			String HOME = SCRIPT_DIR + TEST_DIR;
 			fullDMLScriptName = HOME + TEST_NAME + ".dml";
-			programArgs = new String[] {"-explain", "hops", "-stats", "-args", input("X"), input("Y"), output("R")};
+			programArgs = new String[] {"-explain", "hops", "-stats",
+				"-args", input("X"), input("Y"), output("R")};
 			fullRScriptName = HOME + TEST_NAME + ".R";
 			rCmd = getRCmd(inputDir(), expectedDir());
 
-			OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES = rewrites;
+			OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = rewrites;
 			OptimizerUtils.ALLOW_SUM_PRODUCT_REWRITES = rewrites;
+
 			double[][] X = getRandomMatrix(rows, cols, -1, 1, sparsities[0], 7);
 			double[][] Y = getRandomMatrix(cols, 1, -1, 1, sparsities[1], 3);
 			long X_nnz = Stream.of(X).mapToLong(row -> DoubleStream.of(row).filter(val -> val != 0).count()).sum();
@@ -164,9 +186,19 @@ private void testRewriteMatrixMultChainOpSparse(boolean rewrites) {
 			}
 		}
 		finally {
-			OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES = oldFlag1;
+			OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = oldFlag1;
 			OptimizerUtils.ALLOW_SUM_PRODUCT_REWRITES = oldFlag2;
+			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));
+	}
 }

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

@@ -93,7 +93,7 @@ public void testMMChainFourNoRewrite() {
 
 	private void testMMChainWithTransposeOperator(String testname, int numOptTranspositions,
 		int numOriginalTranspositions, boolean rewrites) {
-		boolean oldFlag = OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES;
+		boolean oldFlag = OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES;
 		try {
 			TestConfiguration config = getTestConfiguration(testname);
 			loadTestConfiguration(config);
@@ -104,7 +104,7 @@ private void testMMChainWithTransposeOperator(String testname, int numOptTranspo
 			fullRScriptName = HOME + testname + ".R";
 			rCmd = getRCmd(inputDir(), expectedDir());
 
-			OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES = rewrites;
+			OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = rewrites;
 
 			//execute tests
 			runTest(true, false, null, -1);
@@ -124,7 +124,7 @@ private void testMMChainWithTransposeOperator(String testname, int numOptTranspo
 
 		}
 		finally {
-			OptimizerUtils.ALLOW_ADVANCED_MMCHAIN_REWRITES = oldFlag;
+			OptimizerUtils.ALLOW_TRANSPOSE_MMCHAIN_REWRITES = oldFlag;
 			Recompiler.reinitRecompiler();
 		}
 	}