[SYSTEMDS-3909] New einsum expression evaluation framework

Closes #2312.
Closes #2265.

Author: HubertKrawczyk

src/main/java/org/apache/sysds/common/Builtins.java

@@ -403,6 +403,7 @@ public enum Builtins {
 	UNDER_SAMPLING("underSampling", true),
 	UNIQUE("unique", false, true),
 	UPPER_TRI("upper.tri", false, true),
+	EINSUM("einsum", false, false),
 	XDUMMY1("xdummy1", true), //error handling test
 	XDUMMY2("xdummy2", true); //error handling test
 

src/main/java/org/apache/sysds/common/InstructionType.java

@@ -62,6 +62,7 @@ public enum InstructionType {
 	PMMJ,
 	MMChain,
 	Union,
+	EINSUM,
 
 	//SP Types
 	MAPMM,

src/main/java/org/apache/sysds/common/Opcodes.java

@@ -174,6 +174,7 @@ public enum Opcodes {
 	RBIND("rbind", InstructionType.BuiltinNary),
 	EVAL("eval", InstructionType.BuiltinNary),
 	LIST("list", InstructionType.BuiltinNary),
+	EINSUM("einsum", InstructionType.BuiltinNary),
 
 	//Parametrized builtin functions
 	AUTODIFF("autoDiff", InstructionType.ParameterizedBuiltin),

src/main/java/org/apache/sysds/common/Types.java

@@ -767,8 +767,8 @@ public String toString() {
 
 	/** Operations that require a variable number of operands*/
 	public enum OpOpN {
-		PRINTF, CBIND, RBIND, MIN, MAX, PLUS, MULT, EVAL, LIST;
-		
+		PRINTF, CBIND, RBIND, MIN, MAX, PLUS, MULT, EVAL, LIST, EINSUM;
+
 		public boolean isCellOp() {
 			return this == MIN || this == MAX || this == PLUS || this == MULT;
 		}

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

@@ -26,6 +26,7 @@
 import org.apache.sysds.lops.Lop;
 import org.apache.sysds.common.Types.ExecType;
 import org.apache.sysds.lops.Nary;
+import org.apache.sysds.runtime.einsum.EinsumEquationValidator;
 import org.apache.sysds.runtime.meta.DataCharacteristics;
 import org.apache.sysds.runtime.meta.MatrixCharacteristics;
 
@@ -235,6 +236,14 @@ public void refreshSizeInformation() {
 				setDim1(getInput().size());
 				setDim2(1);
 				break;
+			case EINSUM:
+				String equationString = ((LiteralOp) _input.get(0)).getStringValue();
+				var dims = EinsumEquationValidator.validateEinsumEquationAndReturnDimensions(equationString, this.getInput().subList(1, this.getInput().size()));
+
+				setDim1(dims.getLeft());
+				setDim2(dims.getMiddle());
+				setDataType(dims.getRight());
+				break;
 			case PRINTF:
 			case EVAL:
 				//do nothing:

src/main/java/org/apache/sysds/hops/codegen/cplan/CNodeCell.java

@@ -31,7 +31,7 @@
 
 public class CNodeCell extends CNodeTpl 
 {
-	protected static final String JAVA_TEMPLATE = 
+	public static final String JAVA_TEMPLATE =
 		  "package codegen;\n"
 		+ "import org.apache.sysds.runtime.codegen.LibSpoofPrimitives;\n"
 		+ "import org.apache.sysds.runtime.codegen.SpoofCellwise;\n"

src/main/java/org/apache/sysds/hops/codegen/cplan/CNodeData.java

@@ -56,6 +56,14 @@ public CNodeData(CNodeData node, String newName) {
 		_cols = node.getNumCols();
 		_dataType = node.getDataType();
 	}
+
+	public CNodeData(String name, long hopID, long rows, long cols, DataType dataType) {
+		_name = name;
+		_hopID = hopID;
+		_rows = rows;
+		_cols = cols;
+		_dataType = dataType;
+	}
 
 	@Override
 	public String getVarname() {

src/main/java/org/apache/sysds/lops/Nary.java

@@ -111,6 +111,7 @@ private String getOpcode() {
 			case RBIND:
 			case EVAL:
 			case LIST:
+			case EINSUM:
 				return operationType.name().toLowerCase();
 			case MIN:
 			case MAX:

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

@@ -23,6 +23,7 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedList;
 
 import org.antlr.v4.runtime.ParserRuleContext;
 import org.apache.commons.lang3.ArrayUtils;
@@ -35,6 +36,7 @@
 import org.apache.sysds.conf.ConfigurationManager;
 import org.apache.sysds.hops.OptimizerUtils;
 import org.apache.sysds.parser.LanguageException.LanguageErrorCodes;
+import org.apache.sysds.runtime.einsum.EinsumEquationValidator;
 import org.apache.sysds.runtime.meta.MatrixCharacteristics;
 import org.apache.sysds.runtime.util.DnnUtils;
 import org.apache.sysds.runtime.util.UtilFunctions;
@@ -751,7 +753,9 @@ else if(((ConstIdentifier) getThirdExpr().getOutput())
 			else
 				raiseValidateError("Compress/DeCompress instruction not allowed in dml script");
 			break;
-							
+		case EINSUM:
+			validateEinsum((DataIdentifier) getOutputs()[0]);
+			break;
 		default: //always unconditional
 			raiseValidateError("Unknown Builtin Function opcode: " + _opcode, false);
 		}
@@ -2063,7 +2067,9 @@ else if(this.getOpCode() == Builtins.MAX_POOL || this.getOpCode() == Builtins.AV
 			output.setValueType(ValueType.INT64);
 			output.setNnz(id.getDim2());
 			break;
-
+		case EINSUM:
+			validateEinsum(output);
+			break;
 		default:
 			if( isMathFunction() ) {
 				checkMathFunctionParam();
@@ -2096,6 +2102,49 @@ else if(this.getOpCode() == Builtins.MAX_POOL || this.getOpCode() == Builtins.AV
 		}
 	}
 
+	private void validateEinsum(DataIdentifier output){
+		if(getSecondExpr() == null)
+			raiseValidateError("Einsum: at least one input matrix required", false,
+					LanguageErrorCodes.INVALID_PARAMETERS);
+
+		if(!(getFirstExpr() instanceof StringIdentifier))
+			raiseValidateError("Einsum: first argument has to be equation str", false,
+					LanguageErrorCodes.INVALID_PARAMETERS);
+
+		String equationString = ((StringIdentifier)getFirstExpr()).getValue();
+
+		if (equationString.length() == 0) raiseValidateError("Einsum: equation str too short", false, LanguageErrorCodes.INVALID_PARAMETERS);
+		if (equationString.charAt(0) == '-' || equationString.charAt(0) == ',') raiseValidateError("Einsum: equation str invalid", false, LanguageErrorCodes.INVALID_PARAMETERS);
+
+		Expression[] expressions = getAllExpr();
+		boolean allDimsKnown = true;
+
+		LinkedList<Identifier> matrixBlocks = new LinkedList<>();
+		for (int i=1;i<expressions.length; i++){
+			checkMatrixParam(expressions[i]);
+			if(!(expressions[i]).getOutput().dimsKnown()){
+				allDimsKnown = false;
+				break;
+			}
+			matrixBlocks.add((expressions[i].getOutput()));
+		}
+
+		if(allDimsKnown){
+			var dims = EinsumEquationValidator.validateEinsumEquationAndReturnDimensions(equationString, matrixBlocks);
+
+			output.setDataType(dims.getRight());
+			output.setDimensions(dims.getLeft(), dims.getMiddle());
+		}else{
+			DataType dataType = EinsumEquationValidator.validateEinsumEquationNoDimensions(equationString, _args.length - 1);
+
+			output.setDataType(dataType);
+			output.setDimensions(-1l, -1l);
+		}
+
+		output.setValueType(ValueType.FP64);
+		output.setBlocksize(getSecondExpr().getOutput().getBlocksize());
+	}
+
 	private void setBinaryOutputProperties(DataIdentifier output) {
 		DataType dt1 = getFirstExpr().getOutput().getDataType();
 		DataType dt2 = getSecondExpr().getOutput().getDataType();

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

@@ -2447,7 +2447,10 @@ private Hop processBuiltinFunctionExpression(BuiltinFunctionExpression source, D
 				new NaryOp(target.getName(), target.getDataType(), target.getValueType(),
 					OpOpN.valueOf(source.getOpCode().name()), processAllExpressions(source.getAllExpr(), hops));
 			break;
-
+		case EINSUM:
+			currBuiltinOp = new NaryOp(target.getName(), target.getDataType(), target.getValueType(),
+					OpOpN.valueOf(source.getOpCode().name()), processAllExpressions(source.getAllExpr(), hops));
+			break;
 		case PPRED:
 			String sop = ((StringIdentifier)source.getThirdExpr()).getValue();
 			sop = sop.replace("\"", "");