BatchMatMulFusionPass

Description

There are two fusion modes:

Mode 1: The transpose nodes that meet the following pattern relationship are eliminated.

Before: After:

Mode 2: When the input shape of BatchMatMul is 2D, BatchMatMul is converted to MatMul. This fusion pattern is disabled by default.

Restrictions

The restrictions on mode 1 are as follows:

  • Transpose1 and Transpose2 can coexist, or only one of them exists.
  • Transpose types include Transpose and TransposeD.
  • The Transpose node swaps only the last two dimensions of the input. For example, if the input shape of the Transpose node is (batch, a, b), the output shape is (batch, b, a).
  • The matmul types include BatchMatMul/BatchMatMulV2/MatMul/MatMulV2.
  • The input dtype of the matmul node can only be float16, float32, or bfloat16.