BatchMatMul2TransposeBatchMatMulFusionPass

Description

Mode 1

When the Transpose (optional), BatchMatMul, and Transpose operators are connected in the sequence shown in the following figure, they can be fused into the TransposeBatchMatmul operator.

For Atlas 350 Accelerator Card, the fusion pattern is:

Mode 2

When the shape of the x1 input of the BatchMatMul/BatchMatMulV2 node is [B2,B1,1,K] and the shape of the x2 input is [1,B1,K,N] or [1,B1,N,K], a Reshape node is inserted to merge axes, resulting in x1 [B2,B1,K] and x2 [B1,K,N] or [B1,N,K]. Then a TransposeBatchMatMul node is called to produce the result [B2,B1,N]. Another Reshape node is inserted to transform the result into [B2,B1,1,N].

Mode 3

When the Transpose (optional), Transpose (optional), BatchMatMul, Transpose, Reshape, Reshape, and Transpose operators are connected in the sequence shown in the following figure, they can be fused into the TransposeBatchMatmul operator.

Restrictions

Mode 1

• The supported data types of the inputs x1 and x2 and the output out (the inputs must correspond to the output) are BFLOAT16, FLOAT16, and FLOAT32.
• The input and output data formats must be ND.
• Input x1 supports [B,M,K] or [M,B,K], and input x2 supports only [B,K,N]. The corresponding attributes of TransposeBatchMatMul are perm_x1=[0,1,2]/[1,0,2] and perm_x2=[0,1,2]. This restriction applies only to the following models:
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas A3 training product/Atlas A3 inference product
• Input x1 is [B,M,K] or [M,B,K], and input x2 is [B,K,N] or [B,N,K]. The corresponding attributes of TransposeBatchMatMul are perm_x1=[0,1,2]/[1,0,2] and perm_x2=[0,1,2]/[0,2,1]. This restriction applies to the following models:
  • Atlas 350 Accelerator Card

Mode 2

• Input x1 is [B2,B1,1,K] and input x2 is [1,B1,K,N] or [1,B1,N,K]. The corresponding BatchMatmul attributes are adj_x1=false and adj_x2=false/true.
• The supported data types of the inputs x1 and x2 and the output out (the inputs must correspond to the output) are FLOAT32.
• The input and output data formats must be ND.
• This fusion pattern does not support HFLOAT32. This restriction applies only to the following models:
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas A3 training product/Atlas A3 inference product
• The input must meet the following condition: B1 × K < 65536. This restriction applies only to the following models:
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas A3 training product/Atlas A3 inference product
• When the input data type is BFLOAT16 or FLOAT16, K and N are 128-aligned. The following conditions must be met: B × K < 65536 or B × K ≥ 65536 (K < 65536). This restriction applies only to the following models:
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas A3 training product/Atlas A3 inference product
• When the input data type is FLOAT32 and the input Transpose nodes are not empty, there is no alignment restriction, but the following condition must be met: B × K < 65536. This restriction applies only to the following models:
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas A3 training product/Atlas A3 inference product

Mode 3

• Output out is [B1,M,B/B1*N], where B must be divisible by B1.
• The supported data types of the inputs x1 and x2 and the output out (the inputs must correspond to the output) are BFLOAT16, FLOAT16, and FLOAT32.
• The input and output data formats must be ND.
• Input x1 supports [B,M,K] or [M,B,K], and input x2 supports only [B,K,N]. The corresponding attributes of TransposeBatchMatMul are perm_x1=[0,1,2]/[1,0,2] and perm_x2=[0,1,2]. This restriction applies only to the following models:
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas A3 training product/Atlas A3 inference product
• Input x1 is [B,M,K] or [M,B,K], and input x2 is [B,K,N] or [B,N,K]. The corresponding attributes of TransposeBatchMatMul are perm_x1=[0,1,2]/[1,0,2] and perm_x2=[0,1,2]/[0,2,1]. This restriction applies only to the following models:

  Atlas 350 Accelerator Card

• When the input data type is BFLOAT16 or FLOAT16, K and N are 128-aligned. The following conditions must be met: B × K < 65536 or B × K ≥ 65536 (K < 65536). This restriction applies only to the following models:
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas A3 training product/Atlas A3 inference product
• When the input data type is FLOAT32 and the input Transpose nodes are not empty, there is no alignment restriction, but the following condition must be met: B × K < 65536. This restriction applies only to the following models:
  • Atlas A2 training product/Atlas A2 inference product
  • Atlas A3 training product/Atlas A3 inference product

Availability

Atlas A2 training product/Atlas A2 inference product

Atlas A3 training product/Atlas A3 inference product

Atlas 350 Accelerator Card