[SYSTEMDS-3798] Fix generality of loop vectorization rewrite

Author: mboehm7

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

@@ -138,13 +138,11 @@ private static StatementBlock vectorizeScalarAggregate( StatementBlock sb, State
 					&& right.getInput(0) instanceof IndexingOp )
 				{
 					IndexingOp ix = (IndexingOp)right.getInput(0);
-					if( ix.isRowLowerEqualsUpper() && ix.getInput(1) instanceof DataOp
-						&& ix.getInput(1).getName().equals(itervar) ){
+					if( checkItervarIndexing(ix, itervar, true) ){
 						leftScalar = true;
 						rowIx = true;
 					}
-					else if( ix.isColLowerEqualsUpper() && ix.getInput(3) instanceof DataOp
-						&& ix.getInput(3).getName().equals(itervar) ){
+					else if( checkItervarIndexing(ix, itervar, false) ){
 						leftScalar = true;
 						rowIx = false;
 					}
@@ -157,13 +155,11 @@ else if( HopRewriteUtils.isValidOp(bop.getOp(), MAP_SCALAR_AGGREGATE_SOURCE_OPS)
 					&& left.getInput(0) instanceof IndexingOp )
 				{
 					IndexingOp ix = (IndexingOp)left.getInput(0);
-					if( ix.isRowLowerEqualsUpper() && ix.getInput(1) instanceof DataOp
-						&& ix.getInput(1).getName().equals(itervar) ){
+					if( checkItervarIndexing(ix, itervar, true) ){
 						rightScalar = true;
 						rowIx = true;
 					}
-					else if( ix.isColLowerEqualsUpper() && ix.getInput(3) instanceof DataOp
-						&& ix.getInput(3).getName().equals(itervar) ){
+					else if( checkItervarIndexing(ix, itervar, false) ){
 						rightScalar = true;
 						rowIx = false;
 					}
@@ -231,17 +227,29 @@ private static StatementBlock vectorizeScalarAggregate2( StatementBlock sb, Stat
 					&& root.getName().equals(left.getName())
 					&& right instanceof IndexingOp && right.isScalar())
 				{
-					leftScalar = true;
-					rowIx = true; //row and col
+					if( checkItervarIndexing((IndexingOp)right, itervar, true) ){
+						leftScalar = true;
+						rowIx = true;
+					}
+					else if( checkItervarIndexing((IndexingOp)right, itervar, false) ){
+						leftScalar = true;
+						rowIx = false;
+					}
 				}
 				//check for right scalar plus
 				else if( HopRewriteUtils.isValidOp(bop.getOp(), MAP_SCALAR_AGGREGATE_SOURCE_OPS)  
 					&& right instanceof DataOp && right.getDataType() == DataType.SCALAR
 					&& root.getName().equals(right.getName()) 
 					&& left instanceof IndexingOp && left.isScalar())
 				{
-					rightScalar = true;
-					rowIx = true; //row and col
+					if( checkItervarIndexing((IndexingOp)left, itervar, true) ){
+						rightScalar = true;
+						rowIx = true;
+					}
+					else if( checkItervarIndexing((IndexingOp)left, itervar, false) ){
+						rightScalar = true;
+						rowIx = false;
+					}
 				}
 			}
 		}
@@ -461,6 +469,12 @@ private static StatementBlock vectorizeIndexedCopy( StatementBlock sb, Statement
 		return ret;
 	}
 
+	private static boolean checkItervarIndexing(IndexingOp ix, String itervar, boolean row) {
+		return ix.isRowLowerEqualsUpper() 
+			&& ix.getInput(row?1:3) instanceof DataOp
+			&& ix.getInput(row?1:3).getName().equals(itervar);
+	}
+	
 	private static boolean[] checkLeftAndRightIndexing(LeftIndexingOp lix, IndexingOp rix, String itervar) {
 		boolean[] ret = new boolean[2]; //apply, rowIx
 

src/test/java/org/apache/sysds/test/functions/vect/AutoVectorizationTest.java

@@ -213,42 +213,36 @@ public void testVectorizeForLoopBinaryColNeg() {
 		runVectorizationTest( TEST_NAME24 ); 
 	}
 
-	/**
-	 * 
-	 * @param cfc
-	 * @param vt
-	 */
 	private void runVectorizationTest( String testName ) 
 	{
 		String TEST_NAME = testName;
 
 		try
-		{		
+		{
 			TestConfiguration config = getTestConfiguration(TEST_NAME);
 			loadTestConfiguration(config);
 
-		    String HOME = SCRIPT_DIR + TEST_DIR;
+			String HOME = SCRIPT_DIR + TEST_DIR;
 			fullDMLScriptName = HOME + TEST_NAME + ".dml";
 			programArgs = new String[]{"-explain","-args", input("A"), output("R") };
 
 			fullRScriptName = HOME + TEST_NAME + ".R";
-			rCmd = getRCmd(inputDir(), expectedDir());		
+			rCmd = getRCmd(inputDir(), expectedDir());
 
 			//generate input
 			double[][] A = getRandomMatrix(rows, cols, 0, 1, 1.0, 7);
 			writeInputMatrixWithMTD("A", A, true);	
 
 			//run tests
-	        runTest(true, false, null, -1);
-	        runRScript(true);
-	        
-	        //compare results
-	        HashMap<CellIndex, Double> dmlfile = readDMLMatrixFromOutputDir("R");
+			runTest(true, false, null, -1);
+			runRScript(true);
+			
+			//compare results
+			HashMap<CellIndex, Double> dmlfile = readDMLMatrixFromOutputDir("R");
 			HashMap<CellIndex, Double> rfile  = readRMatrixFromExpectedDir("R");
 			TestUtils.compareMatrices(dmlfile, rfile, 1e-14, "DML", "R");		
 		}
-		catch(Exception ex)
-		{
+		catch(Exception ex) {
 			throw new RuntimeException(ex);
 		}
 	}