Remove the pattern unstick_4ds_squeeze_stick_3ds (#3062)

Signed-off-by: Tung D. Le <[email protected]>

Author: tungld

src/Accelerators/NNPA/Dialect/ZHigh/ZHighOps/OpHelper.cpp

@@ -636,5 +636,25 @@ IntegerAttr getDefaultSaturation(PatternRewriter &rewriter) {
     return IntegerAttr();
 }
 
+// Create an array tensor to contain three dimensions of layout 3DS.
+// The tensor is created from 4DS's shape by removing the value 1 at axis 1.
+// e.g. 4DS tensor: tensor<3, 1, 4, 5>,
+// this function returns a tensor: tensor<3xi64> = [3, 4, 5]
+Value create3DSShapeFrom4DS(OpBuilder &builder, Location loc, Value val4DS) {
+  OnnxBuilder create(builder, loc);
+  ArrayRef<int64_t> shape4DS = getShape(val4DS.getType());
+  assert(shape4DS.size() == 4 && "The tensor must have rank of 4");
+  assert(shape4DS[1] == 1 && "The second dim must be 1");
+  if (hasStaticShape(val4DS.getType())) {
+    return create.constantInt64(
+        ArrayRef<int64_t>{shape4DS[0], shape4DS[2], shape4DS[3]});
+  }
+  Value dim0 = create.dim(val4DS, 0);
+  Value dim1 = create.dim(val4DS, 2);
+  Value dim2 = create.dim(val4DS, 3);
+  return create.concat(
+      RankedTensorType::get({3}, builder.getI64Type()), {dim0, dim1, dim2}, 0);
+}
+
 } // namespace zhigh
 } // namespace onnx_mlir

src/Accelerators/NNPA/Dialect/ZHigh/ZHighOps/OpHelper.hpp

@@ -110,6 +110,13 @@ bool hasNNPAUse(mlir::Value v);
 /// Get saturation settings.
 mlir::IntegerAttr getDefaultSaturation(mlir::PatternRewriter &rewriter);
 
+/// Create an array tensor to contain three dimensions of layout 3DS.
+/// The tensor is created from 4DS's shape by removing the value 1 at axis 1.
+/// e.g. 4DS tensor: tensor<3, 1, 4, 5>,
+/// this function returns a tensor: tensor<3xi64> = [3, 4, 5]
+mlir::Value create3DSShapeFrom4DS(
+    mlir::OpBuilder &builder, mlir::Location loc, mlir::Value val3DS);
+
 } // namespace zhigh
 } // namespace onnx_mlir
 #endif

src/Accelerators/NNPA/Dialect/ZHigh/ZHighOps/OpHelper.td

@@ -109,6 +109,13 @@ def TensorHas3DSLayout: Constraint<
   "ZTensor has 3DS layout"
 >;
 
+def TensorHas4DSLayout: Constraint<
+  CPred<"::onnx_mlir::is4DSLayout("
+        "::onnx_mlir::zhigh::convertZTensorDataLayoutToStringAttr($_builder, "
+        "::onnx_mlir::zhigh::getZTensorLayout($0.getType())))">,
+  "ZTensor has 4DS layout"
+>;
+
 def TensorHasNHWCLayout: Constraint<
   CPred<"::onnx_mlir::isNHWCLayout("
         "::onnx_mlir::zhigh::convertZTensorDataLayoutToStringAttr($_builder, "
@@ -238,4 +245,17 @@ def GetDefaultSaturation : NativeCodeCall<
   "::onnx_mlir::zhigh::getDefaultSaturation($_builder)"
 >;
 
+class IsConstOf<int v>: Constraint<
+  CPred<"onnx_mlir::isConstOf($0, " # v # ")">,
+  "Value is a s scalar constant of v"
+>;
+
+// Create an array tensor to contain three dimensions of layout 3DS.
+// The tensor is created from 4DS's shape by removing the value 1 at axis 1.
+// e.g. 4DS tensor: tensor<3, 1, 4, 5>,
+// this function returns a tensor: tensor<3xi64> = [3, 4, 5]
+def Create3DSShapeFrom4DS: NativeCodeCall<
+  "::onnx_mlir::zhigh::create3DSShapeFrom4DS($_builder, $_loc, $0)"
+>;
+
 #endif // OP_HELPER

src/Accelerators/NNPA/Dialect/ZHigh/ZHighOps/Stick/Stick.cpp

@@ -13,6 +13,7 @@
 
 #include "src/Accelerators/NNPA/Dialect/ZHigh/ZHighOps/ShapeHelper.hpp"
 #include "src/Compiler/CompilerOptions.hpp"
+#include "src/Dialect/ONNX/ONNXOps/OpHelper.hpp"
 
 using namespace mlir;
 using namespace onnx_mlir;
@@ -138,6 +139,7 @@ void ZHighStickOp::getCanonicalizationPatterns(
   results.insert<NoneTypeStickRemovalPattern>(context);
   results.insert<StickUnstickSameLayoutRemovalPattern>(context);
   results.insert<StickUnstickDiffLayoutRemovalPattern>(context);
+  results.insert<Stick3DSSqueezeUnstick4DSPattern>(context);
   results.insert<ReplaceONNXLeakyReluPattern>(context);
   results.insert<ReplaceONNXSoftplusPattern>(context);
   results.insert<ReplaceONNXReciprocalSqrtPattern>(context);

src/Accelerators/NNPA/Dialect/ZHigh/ZHighOps/Stick/ZHighStick.td

@@ -317,4 +317,23 @@ def ReshapeTransposeReshape3DSTo2DPattern : Pat<
   ]
 >;
 
+// Pattern in the CCFD model.
+// 4DS and 3DS have exactly same data values when the second dim of 4DS is 1.
+// (The second dim of 4DS indicates unidirectional or bidirectional LSTM/GRU/RNN,
+// in which: 1 means unidirectional, 2 means bidirectional)
+// This rewriting is true no matter dims (except the 2nd dim) are dynamic or static.
+def Stick3DSSqueezeUnstick4DSPattern: Pat<
+  // Input: X -> unstick (4DS) -> Squeeze (axis=1) -> stick (3DS).
+  (ZHighStickOp:$stick
+    (ONNXSqueezeOp (ZHighUnstickOp:$unstick $X), $axes),
+    $_, $_),
+  // Output: initial X value unchanged, but transformed with the new layout.
+  (ZHighReshapeOp $X, (Create3DSShapeFrom4DS $X), (GetLayout $stick)),
+  // Conditions.
+  [(TensorHas4DSLayout $X),     // Input is 4DS.
+   (TensorHas3DSLayout $stick), // Output is 3DS.
+   (IsConstOf<1> $axes),        // squeeze at axis 1.
+  ]
+>;
+
 #endif // STICK_TD

src/Accelerators/NNPA/Support/LayoutHelper.cpp

@@ -81,6 +81,10 @@ bool is4DLayout(StringAttr layout) {
   return (layout && layout.getValue().equals_insensitive(LAYOUT_4D));
 }
 
+bool is4DSLayout(StringAttr layout) {
+  return (layout && layout.getValue().equals_insensitive(LAYOUT_4DS));
+}
+
 bool isNHWCLayout(StringAttr layout) {
   return (layout && layout.getValue().equals_insensitive(LAYOUT_NHWC));
 }

src/Accelerators/NNPA/Support/LayoutHelper.hpp

@@ -47,6 +47,7 @@ zdnn_data_layouts convertLayoutAttrToZDNNDataLayout(
 bool is2DLayout(mlir::StringAttr layout);
 bool is3DSLayout(mlir::StringAttr layout);
 bool is4DLayout(mlir::StringAttr layout);
+bool is4DSLayout(mlir::StringAttr layout);
 bool isNHWCLayout(mlir::StringAttr layout);
 
 mlir::StringAttr getNCHWLayoutAttr(mlir::PatternRewriter &rewriter);

test/mlir/accelerators/nnpa/transform/zhigh-combine.mlir

@@ -528,3 +528,40 @@ func.func @test_Roberta_pattern2_bs1_notmod64(%arg0: tensor<1x384x756xf32>, %arg
 // CHECK:         }
 }
 
+// -----
+
+func.func @replace_unstick_squeeze_stick_static(%arg0: tensor<7x1x128x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>) -> tensor<7x128x200xf16, #zhigh.layout<{dataLayout = "3DS"}>> {
+   %cst1 = onnx.Constant dense<1> : tensor<1xi64>
+   %0 = "zhigh.Unstick"(%arg0) : (tensor<7x1x128x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>) -> tensor<7x1x128x200xf32>
+   %1 = "onnx.Squeeze"(%0, %cst1) : (tensor<7x1x128x200xf32>, tensor<1xi64>) -> tensor<7x128x200xf32>
+   %2 = "zhigh.Stick"(%1) {layout = "3DS"} : (tensor<7x128x200xf32>) -> tensor<7x128x200xf16, #zhigh.layout<{dataLayout = "3DS"}>>
+  "func.return"(%2) : (tensor<7x128x200xf16, #zhigh.layout<{dataLayout = "3DS"}>>) -> ()
+
+// CHECK-LABEL:  func.func @replace_unstick_squeeze_stick_static
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<7x1x128x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>) -> tensor<7x128x200xf16, #zhigh.layout<{dataLayout = "3DS"}>> {
+// CHECK:           [[VAR_0_:%.+]] = onnx.Constant dense<[7, 128, 200]> : tensor<3xi64>
+// CHECK:           [[VAR_1_:%.+]] = "zhigh.Reshape"([[PARAM_0_]], [[VAR_0_]]) {layout = "3DS"} : (tensor<7x1x128x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>, tensor<3xi64>) -> tensor<7x128x200xf16, #zhigh.layout<{dataLayout = "3DS"}>>
+// CHECK:           return [[VAR_1_]] : tensor<7x128x200xf16, #zhigh.layout<{dataLayout = "3DS"}>>
+// CHECK:         }
+}
+
+// -----
+
+func.func @replace_unstick_squeeze_stick_dynamic(%arg0: tensor<?x1x?x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>) -> tensor<?x?x200xf16, #zhigh.layout<{dataLayout = "3DS"}>> {
+   %cst1 = onnx.Constant dense<1> : tensor<1xi64>
+   %0 = "zhigh.Unstick"(%arg0) : (tensor<?x1x?x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>) -> tensor<?x1x?x200xf32>
+   %1 = "onnx.Squeeze"(%0, %cst1) : (tensor<?x1x?x200xf32>, tensor<1xi64>) -> tensor<?x?x200xf32>
+   %2 = "zhigh.Stick"(%1) {layout = "3DS"} : (tensor<?x?x200xf32>) -> tensor<?x?x200xf16, #zhigh.layout<{dataLayout = "3DS"}>>
+  "func.return"(%2) : (tensor<?x?x200xf16, #zhigh.layout<{dataLayout = "3DS"}>>) -> ()
+
+// CHECK-LABEL:  func.func @replace_unstick_squeeze_stick_dynamic
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<?x1x?x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>) -> tensor<?x?x200xf16, #zhigh.layout<{dataLayout = "3DS"}>> {
+// CHECK-DAG:       [[VAR_0_:%.+]] = onnx.Constant dense<200> : tensor<1xi64>
+// CHECK-DAG:       [[VAR_1_:%.+]] = "onnx.Dim"([[PARAM_0_]]) {axis = 0 : si64} : (tensor<?x1x?x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>) -> tensor<1xi64>
+// CHECK-DAG:       [[VAR_2_:%.+]] = "onnx.Dim"([[PARAM_0_]]) {axis = 2 : si64} : (tensor<?x1x?x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>) -> tensor<1xi64>
+// CHECK:           [[VAR_3_:%.+]] = "onnx.Concat"([[VAR_1_]], [[VAR_2_]], [[VAR_0_]]) {axis = 0 : si64} : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<3xi64>
+// CHECK:           [[VAR_4_:%.+]] = "zhigh.Reshape"([[PARAM_0_]], [[VAR_3_]]) {layout = "3DS"} : (tensor<?x1x?x200xf16, #zhigh.layout<{dataLayout = "4DS"}>>, tensor<3xi64>) -> tensor<?x?x200xf16, #zhigh.layout<{dataLayout = "3DS"}>>
+// CHECK:           return [[VAR_4_]] : tensor<?x?x200xf16, #zhigh.layout<{dataLayout = "3DS"}>>
+// CHECK:         }
+}
+