Term OOC covariance TSMM

src/main/java/org/apache/sysds/runtime/functionobjects/COV.java

@@ -63,6 +63,8 @@ private COV() {
 	public Data execute(Data in1, double u, double v, double w2) 
 	{
 		CmCovObject cov1=(CmCovObject) in1;
+		if(w2 == 0)
+			return cov1;
 		if(cov1.isCOVAllZeros())
 		{
 			cov1.w=w2;

src/main/java/org/apache/sysds/runtime/instructions/OOCInstructionParser.java

@@ -29,9 +29,10 @@
 import org.apache.sysds.runtime.instructions.ooc.BinaryOOCInstruction;
 import org.apache.sysds.runtime.instructions.ooc.CSVReblockOOCInstruction;
 import org.apache.sysds.runtime.instructions.ooc.CentralMomentOOCInstruction;
+import org.apache.sysds.runtime.instructions.ooc.CovarianceOOCInstruction;
 import org.apache.sysds.runtime.instructions.ooc.CtableOOCInstruction;
-import org.apache.sysds.runtime.instructions.ooc.IndexingOOCInstruction;
 import org.apache.sysds.runtime.instructions.ooc.DataGenOOCInstruction;
+import org.apache.sysds.runtime.instructions.ooc.IndexingOOCInstruction;
 import org.apache.sysds.runtime.instructions.ooc.OOCInstruction;
 import org.apache.sysds.runtime.instructions.ooc.ParameterizedBuiltinOOCInstruction;
 import org.apache.sysds.runtime.instructions.ooc.ReblockOOCInstruction;
@@ -103,6 +104,8 @@ else if(parts.length == 4)
 				return TeeOOCInstruction.parseInstruction(str);
 			case CentralMoment:
 				return  CentralMomentOOCInstruction.parseInstruction(str);
+			case Covariance:
+				return  CovarianceOOCInstruction.parseInstruction(str);
 			case Ctable:
 				return CtableOOCInstruction.parseInstruction(str);
 			case ParameterizedBuiltin:

src/main/java/org/apache/sysds/runtime/instructions/ooc/CovarianceOOCInstruction.java

@@ -0,0 +1,123 @@
+/*
+ * 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.runtime.instructions.ooc;
+
+import java.util.List;
+
+import org.apache.sysds.common.Opcodes;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
+import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
+import org.apache.sysds.runtime.controlprogram.parfor.LocalTaskQueue;
+import org.apache.sysds.runtime.functionobjects.COV;
+import org.apache.sysds.runtime.instructions.InstructionUtils;
+import org.apache.sysds.runtime.instructions.cp.CPOperand;
+import org.apache.sysds.runtime.instructions.cp.CmCovObject;
+import org.apache.sysds.runtime.instructions.cp.DoubleObject;
+import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
+import org.apache.sysds.runtime.matrix.operators.COVOperator;
+import org.apache.sysds.runtime.meta.DataCharacteristics;
+
+public class CovarianceOOCInstruction extends ComputationOOCInstruction {
+
+	private CovarianceOOCInstruction(COVOperator cov, CPOperand in1, CPOperand in2, CPOperand in3, CPOperand out,
+										String opcode, String str) {
+		super(OOCType.COV, cov, in1, in2, in3, out, opcode, str);
+	}
+
+	public static CovarianceOOCInstruction parseInstruction(String str) {
+		String[] parts = InstructionUtils.getInstructionPartsWithValueType(str);
+		String opcode = parts[0];
+
+		if(!opcode.equalsIgnoreCase(Opcodes.COV.toString()))
+			throw new DMLRuntimeException("CovarianceOOCInstruction.parseInstruction():: Unknown opcode " + opcode);
+
+		// the OOC instruction string matches the Spark format,
+
+		COVOperator cov = new COVOperator(COV.getCOMFnObject());
+		if(parts.length == 4) { // this is the case for unweighted cov.A.B.out
+			CPOperand in1 = new CPOperand(parts[1]);
+			CPOperand in2 = new CPOperand(parts[2]);
+			CPOperand out = new CPOperand(parts[3]);
+			return new CovarianceOOCInstruction(cov, in1, in2, null, out, opcode, str);
+		}
+		else if(parts.length == 5) {// this is the case for weighted cov.A.B.W.out
+			CPOperand in1 = new CPOperand(parts[1]);
+			CPOperand in2 = new CPOperand(parts[2]);
+			CPOperand in3 = new CPOperand(parts[3]);
+			CPOperand out = new CPOperand(parts[4]);
+			return new CovarianceOOCInstruction(cov, in1, in2, in3, out, opcode, str);
+		}
+		else {
+			throw new DMLRuntimeException("Invalid number of arguments in Instruction: " + str);
+		}
+	}
+
+	@Override
+	public void processInstruction(ExecutionContext ec) {
+		COVOperator cov_op = (COVOperator) _optr;
+
+		MatrixObject mo1 = ec.getMatrixObject(input1.getName());
+		MatrixObject mo2 = ec.getMatrixObject(input2.getName());
+
+		OOCStream<IndexedMatrixValue> q1 = mo1.getStreamHandle();
+		OOCStream<IndexedMatrixValue> q2 = mo2.getStreamHandle();
+
+		OOCStream<CmCovObject> covObjs = createWritableStream();
+
+		if(input3 == null) {
+			// unweighted covariance join the two tile streams by block index
+			joinOOC(q1, q2, covObjs,
+					(a, b) -> ((MatrixBlock) a.getValue()).covOperations(cov_op, (MatrixBlock) b.getValue()),
+					IndexedMatrixValue::getIndexes);
+		}
+		else {
+			// weighted covariance additionally join the weights tile stream
+			MatrixObject mo3 = ec.getMatrixObject(input3.getName());
+
+			DataCharacteristics dc1 = ec.getDataCharacteristics(input1.getName());
+			DataCharacteristics dc2 = ec.getDataCharacteristics(input2.getName());
+			DataCharacteristics dc3 = ec.getDataCharacteristics(input3.getName());
+			if(dc1.getBlocksize() != dc2.getBlocksize() || dc1.getBlocksize() != dc3.getBlocksize())
+				throw new DMLRuntimeException("Different block sizes are not yet supported");
+
+			OOCStream<IndexedMatrixValue> q3 = mo3.getStreamHandle();
+
+			joinOOC(List.of(q1, q2, q3), covObjs,
+					tiles -> ((MatrixBlock) tiles.get(0).getValue()).covOperations(cov_op,
+							(MatrixBlock) tiles.get(1).getValue(), (MatrixBlock) tiles.get(2).getValue()),
+					IndexedMatrixValue::getIndexes);
+		}
+
+		try {
+			CmCovObject agg = covObjs.dequeue();
+			CmCovObject next;
+
+			while((next = covObjs.dequeue()) != LocalTaskQueue.NO_MORE_TASKS)
+				agg = (CmCovObject) cov_op.fn.execute(agg, next);
+
+			ec.setScalarOutput(output.getName(), new DoubleObject(agg.getRequiredResult(cov_op)));
+		}
+		catch(Exception ex) {
+			throw new DMLRuntimeException(ex);
+		}
+	}
+}

src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCInstruction.java

@@ -80,7 +80,7 @@ public abstract class OOCInstruction extends Instruction {
 
 	public enum OOCType {
 		Reblock, Tee, Binary, Ternary, Unary, AggregateUnary, AggregateBinary, AggregateTernary, MAPMM, MMTSJ,
-		MAPMMCHAIN, Reorg, CM, Ctable, MatrixIndexing, ParameterizedBuiltin, Rand, Append, Quaternary, Reshape
+		MAPMMCHAIN, Reorg, CM, COV, Ctable, MatrixIndexing, ParameterizedBuiltin, Rand, Append, Quaternary, Reshape
 	}
 
 	protected final OOCInstruction.OOCType _ooctype;

src/main/java/org/apache/sysds/runtime/instructions/ooc/TSMMOOCInstruction.java

@@ -19,9 +19,13 @@
 
 package org.apache.sysds.runtime.instructions.ooc;
 
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
 import org.apache.sysds.common.Opcodes;
 import org.apache.sysds.lops.MMTSJ;
 import org.apache.sysds.lops.MMTSJ.MMTSJType;
+import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
 import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
 import org.apache.sysds.runtime.controlprogram.parfor.LocalTaskQueue;
@@ -30,7 +34,9 @@
 import org.apache.sysds.runtime.instructions.InstructionUtils;
 import org.apache.sysds.runtime.instructions.cp.CPOperand;
 import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
+import org.apache.sysds.runtime.matrix.data.LibMatrixReorg;
 import org.apache.sysds.runtime.matrix.data.MatrixBlock;
+import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
 import org.apache.sysds.runtime.matrix.operators.AggregateBinaryOperator;
 import org.apache.sysds.runtime.matrix.operators.AggregateOperator;
 import org.apache.sysds.runtime.matrix.operators.BinaryOperator;
@@ -48,7 +54,7 @@ public static TSMMOOCInstruction parseInstruction(String str) {
 		String[] parts = InstructionUtils.getInstructionPartsWithValueType(str);
 		InstructionUtils.checkNumFields(parts, 3);
 		String opcode = parts[0];
-		CPOperand in1 = new CPOperand(parts[1]); // the large matrix (streamed), columns <= blocksize
+		CPOperand in1 = new CPOperand(parts[1]); // the large matrix (streamed)
 		CPOperand out = new CPOperand(parts[2]);
 		MMTSJ.MMTSJType mmtsjType = MMTSJ.MMTSJType.valueOf(parts[3]);
 
@@ -59,39 +65,116 @@ public static TSMMOOCInstruction parseInstruction(String str) {
 	}
 
 	@Override
-	public void processInstruction( ExecutionContext ec ) {
+	public void processInstruction(ExecutionContext ec) {
 		MatrixObject min = ec.getMatrixObject(input1);
-		int nRows = (int) min.getDataCharacteristics().getRows();
-		int nCols = (int) min.getDataCharacteristics().getCols();
-		int bLen = min.getDataCharacteristics().getBlocksize();
-
-		OOCStream<IndexedMatrixValue> qIn = min.getStreamHandle();
+		int numRowBlocks = Math.toIntExact(min.getDataCharacteristics().getNumRowBlocks());
+		int numColBlocks = Math.toIntExact(min.getDataCharacteristics().getNumColBlocks());
+		if((_type.isLeft() && numColBlocks == 1) || (_type.isRight() && numRowBlocks == 1)) {
+			processSingleOutputTileInstruction(ec, min);
+			return;
+		}
+
+		int blocksPerJoinGroup = _type.isLeft() ? numColBlocks : numRowBlocks;
+		int partialsPerOutput = _type.isLeft() ? numRowBlocks : numColBlocks;
+
+		OOCStreamable<IndexedMatrixValue> inputStreamable = min.getStreamable();
+		final boolean createdCache = !inputStreamable.hasStreamCache();
+		final CachingStream inputCache = createdCache ? new CachingStream(min.getStreamHandle())
+			: inputStreamable.getStreamCache();
+
+		OOCStream<List<IndexedMatrixValue>> groupedPartials = createWritableStream();
+		OOCStream<IndexedMatrixValue> partials = createWritableStream();
+		OOCStream<IndexedMatrixValue> out = createWritableStream();
+		addOutStream(out);
+		ec.getMatrixObject(output).setStreamHandle(out);
+
+		CompletableFuture<Void> joinFuture = joinManyOOC(inputCache.getReadStream(), inputCache.getReadStream(), groupedPartials,
+			this::createPartialOutputTiles, this::getJoinIndex, this::getJoinIndex,
+			blocksPerJoinGroup, blocksPerJoinGroup);
+		CompletableFuture<Void> expandFuture = expandOOC(groupedPartials, partials, values -> values);
+
 		BinaryOperator plus = InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString());
+		CompletableFuture<Void> outFuture = groupedReduceOOC(partials, out, (left, right) -> {
+			MatrixBlock result = ((MatrixBlock) left.getValue()).binaryOperations(plus, right.getValue());
+			left.setValue(result);
+			return left;
+		}, partialsPerOutput);
 
-		//validation check TODO extend compiler to not create OOC otherwise
-		if(    (_type.isLeft() && nCols > bLen)
-			|| (_type.isRight() && nRows > bLen) )
-		{
-			throw new UnsupportedOperationException();
-		}
-
-		//int dim = _type.isLeft() ? nCols : nRows;
+		propagateFailuresToOutput(out, List.of(joinFuture, expandFuture, outFuture));
+
+		outFuture.whenComplete((result, error) -> {
+			if(createdCache)
+				inputCache.scheduleDeletion();
+		});
+	}
+
+	private void processSingleOutputTileInstruction(ExecutionContext ec, MatrixObject min) {
+		OOCStream<IndexedMatrixValue> qIn = min.getStreamHandle();
+		BinaryOperator plus = InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString());
 		MatrixBlock resultBlock = null;
 
 		OOCStream<MatrixBlock> tmpStream = createWritableStream();
-
 		mapOOC(qIn, tmpStream,
 			tmp -> ((MatrixBlock) tmp.getValue())
 				.transposeSelfMatrixMultOperations(new MatrixBlock(), _type));
 
 		MatrixBlock tmp;
-		while ((tmp = tmpStream.dequeue()) != LocalTaskQueue.NO_MORE_TASKS) {
-			if (resultBlock == null)
+		while((tmp = tmpStream.dequeue()) != LocalTaskQueue.NO_MORE_TASKS) {
+			if(resultBlock == null)
 				resultBlock = tmp;
 			else
 				resultBlock.binaryOperationsInPlace(plus, tmp);
 		}
 
 		ec.setMatrixOutput(output.getName(), resultBlock);
 	}
+
+	private long getJoinIndex(IndexedMatrixValue value) {
+		return _type.isLeft() ? value.getIndexes().getRowIndex() : value.getIndexes().getColumnIndex();
+	}
+
+	private long getOutputIndex(IndexedMatrixValue value) {
+		return _type.isLeft() ? value.getIndexes().getColumnIndex() : value.getIndexes().getRowIndex();
+	}
+
+	private List<IndexedMatrixValue> createPartialOutputTiles(IndexedMatrixValue left, IndexedMatrixValue right) {
+		long leftIndex = getOutputIndex(left);
+		long rightIndex = getOutputIndex(right);
+		if(leftIndex > rightIndex)
+			return List.of();
+
+		MatrixBlock leftBlock = (MatrixBlock) left.getValue();
+		MatrixBlock rightBlock = (MatrixBlock) right.getValue();
+		if(leftIndex == rightIndex) {
+			MatrixBlock diagonal = leftBlock.transposeSelfMatrixMultOperations(new MatrixBlock(), _type);
+			return List.of(new IndexedMatrixValue(new MatrixIndexes(leftIndex, rightIndex), diagonal));
+		}
+
+		MatrixBlock partial = multiplyOffDiagonal(leftBlock, rightBlock);
+		MatrixBlock mirror = LibMatrixReorg.transpose(partial);
+		return List.of(
+			new IndexedMatrixValue(new MatrixIndexes(leftIndex, rightIndex), partial),
+			new IndexedMatrixValue(new MatrixIndexes(rightIndex, leftIndex), mirror));
+	}
+
+	private MatrixBlock multiplyOffDiagonal(MatrixBlock leftBlock, MatrixBlock rightBlock) {
+		if(_type.isLeft()) {
+			MatrixBlock leftTranspose = LibMatrixReorg.transpose(leftBlock);
+			return leftTranspose.aggregateBinaryOperations(leftTranspose, rightBlock, new MatrixBlock(),
+				(AggregateBinaryOperator) _optr);
+		}
+
+		MatrixBlock rightTranspose = LibMatrixReorg.transpose(rightBlock);
+		return leftBlock.aggregateBinaryOperations(leftBlock, rightTranspose, new MatrixBlock(),
+			(AggregateBinaryOperator) _optr);
+	}
+
+	private static void propagateFailuresToOutput(OOCStream<?> out, List<CompletableFuture<Void>> futures) {
+		for(CompletableFuture<Void> future : futures) {
+			future.exceptionally(error -> {
+				out.propagateFailure(DMLRuntimeException.of(error));
+				return null;
+			});
+		}
+	}
 }