[SYSTEMDS-3332] Support for matrix-vector operations in ternary ifelse(X,B,v)

src/main/java/org/apache/sysds/parser/BuiltinFunctionExpression.java

@@ -2170,19 +2170,36 @@ private void setBinaryOutputProperties(DataIdentifier output) {
 	}
 
 	private void setTernaryOutputProperties(DataIdentifier output, boolean conditional) {
-		DataType dt1 = getFirstExpr().getOutput().getDataType();
-		DataType dt2 = getSecondExpr().getOutput().getDataType();
-		DataType dt3 = getThirdExpr().getOutput().getDataType();
-		DataType dtOut = (dt1.isMatrix() || dt2.isMatrix() || dt3.isMatrix()) ?
-			DataType.MATRIX : DataType.SCALAR;
-		if( dt1==DataType.MATRIX && dt2==DataType.MATRIX )
-			checkMatchingDimensions(getFirstExpr(), getSecondExpr(), false, conditional);
-		if( dt1==DataType.MATRIX && dt3==DataType.MATRIX )
-			checkMatchingDimensions(getFirstExpr(), getThirdExpr(), false, conditional);
-		if( dt2==DataType.MATRIX && dt3==DataType.MATRIX )
-			checkMatchingDimensions(getSecondExpr(), getThirdExpr(), false, conditional);
+		Expression expr1 = getFirstExpr();
+		Expression expr2 = getSecondExpr();
+		Expression expr3 = getThirdExpr();
+		DataType dt1 = expr1.getOutput().getDataType();
+		DataType dt2 = expr2.getOutput().getDataType();
+		DataType dt3 = expr3.getOutput().getDataType();
+		final long r1 = expr1.getOutput().getDim1();
+		final long r2 = expr2.getOutput().getDim1();
+		final long r3 = expr3.getOutput().getDim1();
+		final long c1 = expr1.getOutput().getDim2();
+		final long c2 = expr2.getOutput().getDim2();
+		final long c3 = expr3.getOutput().getDim2();
+		final long m = Math.max(Math.max(r1, r2), r3);
+		final long n = Math.max(Math.max(c1, c2), c3);
+
+		boolean unknownDim = (r1 == -1 || r2 == -1 || r3 == -1 || c1 == -1 || c2 == -1 || c3 == -1);
+		if (!unknownDim && ((dt1 == DataType.MATRIX && r1 != 1 && r1 != m)
+			|| (dt2 == DataType.MATRIX && r2 != 1 && r2 != m)
+			|| (dt3 == DataType.MATRIX && r3 != 1 && r3 != m)
+			|| (dt1 == DataType.MATRIX && c1 != 1 && c1 != n)
+			|| (dt2 == DataType.MATRIX && c2 != 1 && c2 != n)
+			|| (dt3 == DataType.MATRIX && c3 != 1 && c3 != n))) {
+			raiseValidateError("Mismatch in matrix dimensions of parameters for function "
+					+ this.getOpCode(), conditional, LanguageErrorCodes.INVALID_PARAMETERS);
+		}
+
 		MatrixCharacteristics dims1 = getBinaryMatrixCharacteristics(getFirstExpr(), getSecondExpr());
 		MatrixCharacteristics dims2 = getBinaryMatrixCharacteristics(getSecondExpr(), getThirdExpr());
+		DataType dtOut = (dt1.isMatrix() || dt2.isMatrix() || dt3.isMatrix()) ?
+			DataType.MATRIX : DataType.SCALAR;
 		output.setDataType(dtOut);
 		output.setValueType(dtOut==DataType.MATRIX ? ValueType.FP64 :
 			computeValueType(getSecondExpr(), getThirdExpr(), true));

src/main/java/org/apache/sysds/runtime/compress/lib/CLALibTernaryOp.java

@@ -59,7 +59,7 @@ public static MatrixBlock ternaryOperations(TernaryOperator op, MatrixBlock m1,
 		final int n = Math.max(Math.max(c1, c2), c3);
 
 		// double check that the dimensions are valid.
-		MatrixBlock.ternaryOperationCheck(s1, s2, s3, m, r1, r2, r3, n, c1, c2, c3);
+		MatrixBlock.ternaryOperationCheck(op, s1, s2, s3, m, r1, r2, r3, n, c1, c2, c3);
 
 		MatrixBlock ret;
 

src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixTercell.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.Future;
 
 import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.functionobjects.IfElse;
 import org.apache.sysds.runtime.matrix.operators.TernaryOperator;
 import org.apache.sysds.runtime.util.CommonThreadPool;
 import org.apache.sysds.runtime.util.UtilFunctions;
@@ -88,19 +89,87 @@ public static void tercellOp(MatrixBlock m1, MatrixBlock m2, MatrixBlock m3, Mat
 
 	private static long unsafeTernary(MatrixBlock m1, MatrixBlock m2, MatrixBlock m3, MatrixBlock ret,
 		TernaryOperator op, boolean s1, boolean s2, boolean s3, double d1, double d2, double d3, int rl, int ru)
+	{
+		if(op.fn instanceof IfElse) {
+			return unsafeTernaryIfElse(m1, m2, m3, ret, op, s1, s2,
+			s3, d1, d2, d3, rl, ru);
+		}
+		else {
+			return unsafeTernaryDefault(m1, m2, m3, ret, op, s1, s2,
+				s3, d1, d2, d3, rl, ru);
+		}
+	}
+
+	private static long unsafeTernaryIfElse(MatrixBlock m1, MatrixBlock m2, MatrixBlock m3, MatrixBlock ret,
+		TernaryOperator op, boolean s1, boolean s2, boolean s3, double d1, double d2, double d3, int rl, int ru)
+	{
+		// IfElse specific optimizations are applied here
+		int n = ret.clen;
+		long lnnz = 0;
+		final int r1 = m1.getNumRows();
+		final int c1 = m1.getNumColumns();
+		if (c1 == 1) {
+			for( int i=rl; i<ru; i++ ) {
+				double in1 = s1 ? d1 : m1.get(Math.min(i, r1-1), 0);
+				MatrixBlock tmp = (in1 != 0) ? m2 : m3;
+				final int tmpCol = tmp.getNumColumns();
+				final int tmpRow = tmp.getNumRows();
+				final int tmpi = Math.min(i, tmpRow-1);
+				for( int j=0; j<n; j++ ) {
+					double val = tmp.get(tmpi, Math.min(j, tmpCol-1));
+					lnnz += (val != 0) ? 1 : 0;
+					ret.appendValuePlain(i, j, val);
+				}
+			}
+			return lnnz;
+		}
+		else if (r1 == 1) {
+			for( int j=0; j<n; j++ ) {
+				double in1 = s1 ? d1 : m1.get(0, Math.min(j, c1-1));
+				MatrixBlock tmp = (in1 != 0) ? m2 : m3;
+				final int tmpCol = tmp.getNumColumns();
+				final int tmpRow = tmp.getNumRows();
+				final int tmpj = Math.min(j, tmpCol-1);
+				for( int i=rl; i<ru; i++ ) {
+					double val = tmp.get(Math.min(i, tmpRow-1), tmpj);
+					lnnz += (val != 0) ? 1 : 0;
+					ret.appendValuePlain(i, j, val);
+				}
+			}
+			return lnnz;
+		}
+		else {
+			return unsafeTernaryDefault(m1, m2, m3, ret, op, s1, s2,
+				s3, d1, d2, d3, rl, ru);
+		}
+	}
+
+	private static long unsafeTernaryDefault(MatrixBlock m1, MatrixBlock m2, MatrixBlock m3, MatrixBlock ret,
+		TernaryOperator op, boolean s1, boolean s2, boolean s3, double d1, double d2, double d3, int rl, int ru)
 	{
 		//basic ternary operations (all combinations sparse/dense)
 		int n = ret.clen;
 		long lnnz = 0;
-		for( int i=rl; i<ru; i++ )
+		final int r1 = m1.getNumRows();
+		final int r2 = m2.getNumRows();
+		final int r3 = m3.getNumRows();
+		final int c1 = m1.getNumColumns();
+		final int c2 = m2.getNumColumns();
+		final int c3 = m3.getNumColumns();
+
+		for( int i=rl; i<ru; i++ ) {
+			final int i1 = Math.min(i, r1-1);
+			final int i2 = Math.min(i, r2-1);
+			final int i3 = Math.min(i, r3-1);
 			for( int j=0; j<n; j++ ) {
-				double in1 = s1 ? d1 : m1.get(i, j);
-				double in2 = s2 ? d2 : m2.get(i, j);
-				double in3 = s3 ? d3 : m3.get(i, j);
+				double in1 = s1 ? d1 : m1.get(i1, Math.min(j, c1-1));
+				double in2 = s2 ? d2 : m2.get(i2, Math.min(j, c2-1));
+				double in3 = s3 ? d3 : m3.get(i3, Math.min(j, c3-1));
 				double val = op.fn.execute(in1, in2, in3);
 				lnnz += (val != 0) ? 1 : 0;
 				ret.appendValuePlain(i, j, val);
 			}
+		}
 
 		//set global output nnz once
 		return lnnz;

src/main/java/org/apache/sysds/runtime/matrix/data/MatrixBlock.java

@@ -3080,29 +3080,11 @@ public MatrixBlock ternaryOperations(TernaryOperator op, MatrixBlock m2, MatrixB
 		final int n = Math.max(Math.max(c1, c2), c3);
 		final long nnz = nonZeros;
 
-		ternaryOperationCheck(s1, s2, s3, m, r1, r2, r3, n, c1, c2, c3);
+		ternaryOperationCheck(op, s1, s2, s3, m, r1, r2, r3, n, c1, c2, c3);
 
 		//prepare result
 		if( op.fn instanceof IfElse && (s1 || nnz==0 || nnz==(long)m*n) ) {
-
-			ret.reset(m, n, false);
-
-			//SPECIAL CASE for shallow-copy if-else
-			boolean expr = s1 ? (d1 != 0) : (nnz==(long)m*n);
-			MatrixBlock tmp = expr ? m2 : m3;
-			if( tmp.rlen==m && tmp.clen==n ) {
-				//shallow copy incl meta data
-				ret.copyShallow(tmp);
-			}
-			else {
-				//fill output with given scalar value
-				double tmpVal = tmp.get(0, 0);
-				if( tmpVal != 0 ) {
-					ret.allocateDenseBlock();
-					ret.denseBlock.set(tmpVal);
-					ret.nonZeros = (long)m * n;
-				}
-			}
+			ternaryIfElseCopy(s1, m, n, d1, nnz, m2, m3, ret);
 		}
 		else{
 			final boolean PM_Or_MM = (op.fn instanceof PlusMultiply || op.fn instanceof MinusMultiply);
@@ -3141,11 +3123,87 @@ else if (s2 != s3 && (op.fn instanceof PlusMultiply || op.fn instanceof MinusMul
 		return ret;
 	}
 
-	public static void ternaryOperationCheck(boolean s1, boolean s2, boolean s3, int m, int r1, int r2, int r3, int n, int c1, int c2, int c3){
+	/**
+	 * Copies and applies broadcasting to either second or third input to IfElse operation.
+	 *
+	 * If the first value of the IfElse-operation meets certain conditions, then
+	 * either the m2 or m3 matrix inputs of IfElse is the result of the operation.
+	 * This methods handles the copying if the conditions are met.
+	 *
+	 * @param s1  Flag, whether the first matrix is a scalar.
+	 * @param m   Rows of the resulting matrix.
+	 * @param n   Columns of the resulting matrix
+	 * @param d1  Value of the entry 0,0 of the first matrix input.
+	 * @param nnz Non-zero entries of the first matrix
+	 * @param m2  Second matrix input of the IfElse operation
+	 * @param m3  Third matrix input of the IfElse operation
+	 * @param ret Result of the operation, where either m2 or m3 is copied into
+	 */
+	private void ternaryIfElseCopy(boolean s1, int m, int n, double d1, long nnz, MatrixBlock m2, MatrixBlock m3, MatrixBlock ret) {
+		ret.reset(m, n, false);
+
+		boolean expr = s1 ? (d1 != 0) : (nnz==(long)m*n);
+		MatrixBlock tmp = expr ? m2 : m3;
+		if (tmp.rlen==m && tmp.clen==n) {
+			//shallow copy incl meta data
+			ret.copyShallow(tmp);
+		}
+		else if (tmp.rlen==m && tmp.clen==1) {
+			ret.allocateDenseBlock();
+			ret.nonZeros = 0;
+			for (int i = 0; i < m; i++) {
+				double tmpVal = tmp.get(i, 0);
+				if (tmpVal != 0) {
+					ret.denseBlock.fillRow(i, tmp.get(i, 0));
+					ret.nonZeros += n;
+				}
+			}
+			ret.examSparsity();
+		}
+		else if (tmp.rlen==1 && tmp.clen==n) {
+			if (tmp.nonZeros != 0) {
+				ret.allocateDenseBlock();
+				ret.nonZeros = 0;
+				double[] tmpArr = new double[n];
+				for (int i = 0; i < n; i++) {
+					tmpArr[i] = tmp.get(0, i);
+				}
+				for (int i = 0; i < m; i++) {
+					ret.denseBlock.set(i, tmpArr);
+					ret.nonZeros += tmp.nonZeros;
+				}
+				ret.examSparsity();
+			}
+		}
+		else {
+			//fill output with given scalar value
+			double tmpVal = tmp.get(0, 0);
+			if (tmpVal != 0) {
+				ret.allocateDenseBlock();
+				ret.denseBlock.set(tmpVal);
+				ret.nonZeros = (long)m * n;
+			}
+		}
+	}
+
+	public static void ternaryOperationCheck(TernaryOperator op, boolean s1, boolean s2, boolean s3, int m, int r1, int r2, int r3, int n, int c1, int c2, int c3){
 		//error handling 
-		if( (!s1 && (r1 != m || c1 != n))
-		|| (!s2 && (r2 != m || c2 != n))
-		|| (!s3 && (r3 != m || c3 != n)) ) {
+		boolean error = false;
+		if (op.fn instanceof IfElse) {
+			error = ((r1 != 1 && r1 != m)
+				|| (r2 != 1 && r2 != m)
+				|| (r3 != 1 && r3 != m)
+				|| (c1 != 1 && c1 != n)
+				|| (c2 != 1 && c2 != n)
+				|| (c3 != 1 && c3 != n));
+		}
+		else {
+			error = ((!s1 && (r1 != m || c1 != n))
+				|| (!s2 && (r2 != m || c2 != n))
+				|| (!s3 && (r3 != m || c3 != n)));
+		}
+
+		if (error) {
 			throw new DMLRuntimeException("Block sizes are not matched for ternary cell operations: "
 			+ r1 + "x" + c1 + " vs " + r2 + "x" + c2 + " vs " + r3 + "x" + c3);
 		}