[SYSTEMDS-3874] Java17 Vectorized LibMM

This commit adds vectorized kernels for matrix multiplication.

the vector API improves performance for single-threaded execution of our
AMD box improves by ~80% and Intel by ~60% for dense mm. These improvements
are with allocation overhead of the output and in ideal cases where the
input is cached and the JIT compilation is done.

The biggest change for users is that SystemDS now would require
`--add-modules=jdk.incubator.vector` to all execution calls. The commit
appropriately modifies all scripts to do this. However, all calling code
must be modified if it bypasses the bin/systemds and calls Java directly.

To measure the performance difference on your machine, use the added script:

src/test/scripts/performance/matrixMultiplication.sh

Closes #2216

Author: Baunsgaard

bin/systemds

@@ -413,6 +413,7 @@ if [ $WORKER == 1 ]; then
   print_out "#  starting Federated worker on port $PORT"
   CMD=" \
   java $SYSTEMDS_STANDALONE_OPTS \
+  --add-modules=jdk.incubator.vector \
   $LOG4JPROPFULL \
   -jar $SYSTEMDS_JAR_FILE \
   -w $PORT \
@@ -422,6 +423,7 @@ elif [ "$FEDMONITORING" == 1 ]; then
   print_out "#  starting Federated backend monitoring on port $PORT"
   CMD=" \
   java $SYSTEMDS_STANDALONE_OPTS \
+  --add-modules=jdk.incubator.vector \
   $LOG4JPROPFULL \
   -jar $SYSTEMDS_JAR_FILE \
   -fedMonitoring $PORT \
@@ -433,6 +435,7 @@ elif [ $SYSDS_DISTRIBUTED == 0 ]; then
   CMD=" \
   java $SYSTEMDS_STANDALONE_OPTS \
   $LOG4JPROPFULL \
+  --add-modules=jdk.incubator.vector \
   -jar $SYSTEMDS_JAR_FILE \
   -f $SCRIPT_FILE \
   -exec $SYSDS_EXEC_MODE \
@@ -442,6 +445,7 @@ else
   print_out "#  Running script $SCRIPT_FILE distributed with opts: $*"
   CMD=" \
   spark-submit $SYSTEMDS_DISTRIBUTED_OPTS \
+  --add-modules=jdk.incubator.vector \
   $SYSTEMDS_JAR_FILE \
   -f $SCRIPT_FILE \
   -exec $SYSDS_EXEC_MODE \

docker/build.sh

@@ -23,7 +23,7 @@
 # Build the docker containers
 
 # The first build is for running systemds through docker.
-docker image build -f docker/sysds.Dockerfile -t apache/systemds:latest .
+# docker image build -f docker/sysds.Dockerfile -t apache/systemds:latest .
 
 # The second build is for testing systemds. This image installs the R dependencies needed to run the tests.
 docker image build -f docker/testsysds.Dockerfile -t apache/systemds:testing-latest .

docker/entrypoint.sh

@@ -30,7 +30,9 @@ export MAVEN_OPTS="-Xmx512m"
 
 log="/tmp/sysdstest.log"
 mvn -ntp -B test-compile 2>&1 | grep -E "BUILD|Total time:|---|Building SystemDS"
-mvn -ntp -B test -D maven.test.skip=false -D automatedtestbase.outputbuffering=true -D test=$1 2>&1 | grep -v "already exists in destination." | tee $log
+mvn -ntp -B test -D maven.test.skip=false -D automatedtestbase.outputbuffering=true -D test=$1 2>&1 \
+	| grep -v "already exists in destination." \
+	| grep -v 'WARNING: Using incubator modules' | tee $log
 
 # Merge Federated test runs.
 [ -f target/jacoco.exec ] && mv target/jacoco.exec target/jacoco_main.exec

pom.xml

@@ -92,6 +92,7 @@
 			--add-opens=java.base/java.lang.ref=ALL-UNNAMED
 			--add-opens=java.base/java.util.concurrent=ALL-UNNAMED
 			--add-opens=java.base/sun.nio.ch=ALL-UNNAMED
+			--add-modules=jdk.incubator.vector
 		</jvm.addopens>
 	</properties>
 
@@ -357,6 +358,9 @@
 					<source>${java.level}</source>
 					<target>${java.level}</target>
 					<release>${java.level}</release>
+					<compilerArgs>
+						<arg>--add-modules=jdk.incubator.vector</arg>
+					</compilerArgs>
 				</configuration>
 			</plugin>
 
@@ -904,6 +908,7 @@
 							<notimestamp>true</notimestamp>
 							<failOnWarnings>false</failOnWarnings>
 							<quiet>true</quiet>
+							<additionalJOption>--add-modules=jdk.incubator.vector</additionalJOption>
 							<skip>${doc.skip}</skip>
 							<show>public</show>
 							<source>${java.level}</source>

src/main/java/org/apache/sysds/runtime/compress/colgroup/ColGroupDDC.java

@@ -26,8 +26,8 @@
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 
-// import jdk.incubator.vector.DoubleVector;
-// import jdk.incubator.vector.VectorSpecies;
+import jdk.incubator.vector.DoubleVector;
+import jdk.incubator.vector.VectorSpecies;
 import org.apache.commons.lang3.NotImplementedException;
 import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.compress.CompressedMatrixBlock;
@@ -75,7 +75,7 @@ public class ColGroupDDC extends APreAgg implements IMapToDataGroup {
 
 	protected final AMapToData _data;
 
-	// static final VectorSpecies<Double> SPECIES = DoubleVector.SPECIES_PREFERRED;
+	static final VectorSpecies<Double> SPECIES = DoubleVector.SPECIES_PREFERRED;
 
 	private ColGroupDDC(IColIndex colIndexes, IDictionary dict, AMapToData data, int[] cachedCounts) {
 		super(colIndexes, dict, cachedCounts);
@@ -625,16 +625,16 @@ private void identityRightDecompressingMult(MatrixBlock right, MatrixBlock ret,
 		final double[] b = right.getDenseBlockValues();
 		final double[] c = ret.getDenseBlockValues();
 		final int jd = right.getNumColumns();
-		final int vLen = 8;
+		final DoubleVector vVec = DoubleVector.zero(SPECIES);
+		final int vLen = SPECIES.length();
 		final int lenJ = cru - crl;
 		final int end = cru - (lenJ % vLen);
 		for(int i = rl; i < ru; i++) {
 			int k = _data.getIndex(i);
 			final int offOut = i * jd + crl;
 			final double aa = 1;
 			final int k_right = _colIndexes.get(k);
-			vectMM(aa, b, c, end, jd, crl, cru, offOut, k_right, vLen);
-
+			vectMM(aa, b, c, end, jd, crl, cru, offOut, k_right, vLen, vVec);
 		}
 	}
 
@@ -644,8 +644,8 @@ private void defaultRightDecompressingMult(MatrixBlock right, MatrixBlock ret, i
 		final double[] c = ret.getDenseBlockValues();
 		final int kd = _colIndexes.size();
 		final int jd = right.getNumColumns();
-		// final DoubleVector vVec = DoubleVector.zero(SPECIES);
-		final int vLen = 8;
+		final DoubleVector vVec = DoubleVector.zero(SPECIES);
+		final int vLen = SPECIES.length();
 
 		final int blkzI = 32;
 		final int blkzK = 24;
@@ -661,32 +661,22 @@ private void defaultRightDecompressingMult(MatrixBlock right, MatrixBlock ret, i
 					for(int k = bk; k < bke; k++) {
 						final double aa = a[offi + k];
 						final int k_right = _colIndexes.get(k);
-						vectMM(aa, b, c, end, jd, crl, cru, offOut, k_right, vLen);
+						vectMM(aa, b, c, end, jd, crl, cru, offOut, k_right, vLen, vVec);
 					}
 				}
 			}
 		}
 	}
 
-	final void vectMM(double aa, double[] b, double[] c, int endT, int jd, int crl, int cru, int offOut, int k,
-		int vLen) {
-		// vVec = vVec.broadcast(aa);
+	final void vectMM(double aa, double[] b, double[] c, int endT, int jd, int crl, int cru, int offOut, int k, int vLen, DoubleVector vVec) {
+		vVec = vVec.broadcast(aa);
 		final int offj = k * jd;
 		final int end = endT + offj;
 		for(int j = offj + crl; j < end; j += vLen, offOut += vLen) {
-			// DoubleVector res = DoubleVector.fromArray(SPECIES, c, offOut);
-			// DoubleVector bVec = DoubleVector.fromArray(SPECIES, b, j);
-			// res = vVec.fma(bVec, res);
-			// res.intoArray(c, offOut);
-
-			c[offOut] += aa * b[j];
-			c[offOut + 1] += aa * b[j + 1];
-			c[offOut + 2] += aa * b[j + 2];
-			c[offOut + 3] += aa * b[j + 3];
-			c[offOut + 4] += aa * b[j + 4];
-			c[offOut + 5] += aa * b[j + 5];
-			c[offOut + 6] += aa * b[j + 6];
-			c[offOut + 7] += aa * b[j + 7];
+			DoubleVector res = DoubleVector.fromArray(SPECIES, c, offOut);
+			DoubleVector bVec = DoubleVector.fromArray(SPECIES, b, j);
+			res = vVec.fma(bVec, res);
+			res.intoArray(c, offOut);
 		}
 		for(int j = end; j < cru + offj; j++, offOut++) {
 			double bb = b[j];

src/main/java/org/apache/sysds/runtime/compress/colgroup/dictionary/MatrixBlockDictionary.java

@@ -27,6 +27,8 @@
 import java.util.Arrays;
 import java.util.Set;
 
+import jdk.incubator.vector.DoubleVector;
+import jdk.incubator.vector.VectorSpecies;
 import org.apache.commons.lang3.NotImplementedException;
 import org.apache.sysds.runtime.compress.DMLCompressionException;
 import org.apache.sysds.runtime.compress.colgroup.indexes.ArrayIndex;
@@ -65,6 +67,8 @@ public class MatrixBlockDictionary extends ADictionary {
 
 	final private MatrixBlock _data;
 
+	static final VectorSpecies<Double> SPECIES = DoubleVector.SPECIES_PREFERRED;
+
 	/**
 	 * Unsafe private constructor that does not check the data validity. USE WITH CAUTION.
 	 * 
@@ -2088,7 +2092,71 @@ private void preaggValuesFromDenseDictDenseAggArray(final int numVals, final ICo
 
 	private void preaggValuesFromDenseDictDenseAggRange(final int numVals, final IColIndex colIndexes, final int s,
 		final int e, final double[] b, final int cut, final double[] ret) {
-		preaggValuesFromDenseDictDenseAggRangeGeneric(numVals, colIndexes, s, e, b, cut, ret);
+		if(colIndexes instanceof RangeIndex) {
+			RangeIndex ri = (RangeIndex) colIndexes;
+			preaggValuesFromDenseDictDenseAggRangeRange(numVals, ri.get(0), ri.get(0) + ri.size(), s, e, b, cut, ret);
+		}
+		else
+			preaggValuesFromDenseDictDenseAggRangeGeneric(numVals, colIndexes, s, e, b, cut, ret);
+	}
+
+	private void preaggValuesFromDenseDictDenseAggRangeRange(final int numVals, final int ls, final int le, final int rs,
+		final int re, final double[] b, final int cut, final double[] ret) {
+		final int cz = le - ls;
+		final int az = re - rs;
+		// final int nCells = numVals * cz;
+		final double[] values = _data.getDenseBlockValues();
+		// Correctly named ikj matrix multiplication .
+
+		final int blkzI = 32;
+		final int blkzK = 24;
+		final int blkzJ = 1024;
+		for(int bi = 0; bi < numVals; bi += blkzI) {
+			final int bie = Math.min(numVals, bi + blkzI);
+			for(int bk = 0; bk < cz; bk += blkzK) {
+				final int bke = Math.min(cz, bk + blkzK);
+				for(int bj = 0; bj < az; bj += blkzJ) {
+					final int bje = Math.min(az, bj + blkzJ);
+					final int sOffT = rs + bj;
+					final int eOffT = rs + bje;
+					preaggValuesFromDenseDictBlockedIKJ(values, b, ret, bi, bk, bj, bie, bke, cz, az, ls, cut, sOffT, eOffT);
+				}
+			}
+		}
+	}
+
+	private static void preaggValuesFromDenseDictBlockedIKJ(double[] a, double[] b, double[] ret, int bi, int bk, int bj,
+		int bie, int bke, int cz, int az, int ls, int cut, int sOffT, int eOffT) {
+		final int vLen = SPECIES.length();
+		final DoubleVector vVec = DoubleVector.zero(SPECIES);
+		final int leftover = (eOffT - sOffT) % vLen; // leftover not vectorized
+		for(int i = bi; i < bie; i++) {
+			final int offI = i * cz;
+			final int offOutT = i * az + bj;
+			for(int k = bk; k < bke; k++) {
+				final int idb = (k + ls) * cut;
+				final int sOff = sOffT + idb;
+				final int eOff = eOffT + idb;
+				final double v = a[offI + k];
+				vecInnerLoop(v, b, ret, offOutT, eOff, sOff, leftover, vLen, vVec);
+			}
+		}
+	}
+
+	private static void vecInnerLoop(final double v, final double[] b, final double[] ret, final int offOutT,
+		final int eOff, final int sOff, final int leftover, final int vLen, DoubleVector vVec) {
+		int offOut = offOutT;
+		vVec = vVec.broadcast(v);
+		final int end = eOff - leftover;
+		for(int j = sOff; j < end; j += vLen, offOut += vLen) {
+			DoubleVector res = DoubleVector.fromArray(SPECIES, ret, offOut);
+			DoubleVector bVec = DoubleVector.fromArray(SPECIES, b, j);
+			vVec.fma(bVec, res).intoArray(ret, offOut);
+		}
+		for(int j = end; j < eOff; j++, offOut++) {
+			ret[offOut] += v * b[j];
+		}
+
 	}
 
 	private void preaggValuesFromDenseDictDenseAggRangeGeneric(final int numVals, final IColIndex colIndexes,