Implement nGraph transformation to decompose Einsum-7 operation

[rkazants]

Ticket: 54320

Description: Currently the plugins do not support Einsum-7 operation. The implemented nGraph transformation decomposes Einsum-7 operation into a sub-graph that contains operations supported by the plugins. The resulted sub-graph can vary too much and depends on Einsum equation. The sub-graph can include MatMul, Reshape, ReduceSum, Unsqueeze, Multiply operations. The quick check of inference performance on the BERT model including 97 Einsum operations proves the functional correctness of the transformation and does not unveil performance degradation. Also, the transformation is tested on 19 layer tests.

Current Limitations:

1. The transformation supports only Einsum equations that do not have the repeated labels in the single subscript and do not have ellipsis label ('...'). For example, Einsum operation with equation=aab,bc->ac is unsupported due to repeated label a in the first input subscript. In the meantime Einsum operation with equation=abc,dcb->ac is supported.
2. Currently it does not implement algorithm to compute (pseudo-)optimal einsum_path for multiple operand case. And in such cases the operands are contracted consequently. The current implementation is sufficient for now because we only have models with two operand Einsum operation.

Details: The main idea of the transformation is better to present using an example about how numpy einsum operation can be decomposed into simple numpy operations. Let us consider the following code:
input1 = np.random.random_integers(10, size=(2,3,4,5)).astype(np.float)
input2 = np.random.random_integers(10, size=(6,4,5,3)).astype(np.float)
ref_result = np.einsum('aecd,bcde->abe', input1, input2)

In the first step let us group dimensions of both operands into three groups:

1. Common dimensions with labels that are met in both input subscripts and in the output subscript. In this case there is just one common dimension with label e.
2. Reduced dimensions with labels that are met in both input subscripts but are not met in the output subscript. In this case there are two dimensions with labels c and d.
3. Separate dimensions (for each operand) with labels that are met in just one input subscript. In this case, dimension with label a is separate for the first operand; dimension with a label b - for the second operand.

Transpose the first operand so that it has the common dimensions first, the separate dimensions after, and the reduced dimensions lastly. Transpose the second operand so that it has the common dimensions first, the reduced dimensions after, and the separate dimensions lastly. So both subscripts look like as follows: eacd - for the first operands, ecdb - for the second operand
input1_grouped = np.transpose(input1, [1, 0, 2, 3])
input2_grouped = np.transpose(input2, [3, 1, 2, 0])

In the second step let us collapse separate dimensions into one dimension using Reshape operation and do the same for reduced dimensions. It is needed to utilize MatMul operation that has requirements for input format. The common dimensions are sort of batch dimensions for MatMul operation.
input1_reshaped = np.reshape(input1_grouped, (3, 2, 20))
input2_reshaped = np.reshape(input2_grouped, (3, 20, 6))

In the third step, perform MatMul operation.
matmul = np.matmul(input1_reshaped, input2_reshaped)
The result of MatMul operation will stay the common and separate dimensions.

In the fourth step, unroll previously collapsed dimensions (the separate dimensions) using Reshape operation. In this case it does not have collapsed the separate dimensions.
matmul_reshaped = np.reshape(matmul, (3, 2, 6))

Finally, it needs to adjust layout specified by the output subcript. The intermediate result has a layout corresponding to eab subscript but the output subscipt is abe. Hence, perform the transpose.
result = np.transpose(matmul_reshaped, [1, 2, 0])

Values of ref_result and result are matched. The nGraph transformation relies on this idea and its modification to avoid extra transpose using transpose attributes in MatMul operations.
Consider IR with Einsum operation and how it is decomposed. The original IR looks as follows:

The resulted IR before constant-folding:

The idea described above is generalized to multiple operand case by computing intermediate output subscript for each pair of operands.

Signed-off-by: Roman Kazantsev [email protected]

[lazarevevgeny]

In general it looks good. I didn't go into details of logic of some parts of this PR but I trust Roman.

But I agree with Ilya's comment about moving some methods to private section of the class instead of passing transformation instance pointer.

[ilyachur]

Please fix comments in the next PR