feat: Match LayerNorm decomposition pattern with transposing original input and axis

Author: ehsan-toosi

src/Compiler/OnnxToMlirPasses.cpp

@@ -32,13 +32,15 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
       opts.enableConvTransposeDecomposeToPhasedConv,
       opts.enableConvTranspose1dDecomposeToPhasedConv));
   if (!opts.disableRecomposeOption)
-    pm.addNestedPass<func::FuncOp>(onnx_mlir::createRecomposeONNXToONNXPass());
+    pm.addNestedPass<func::FuncOp>(onnx_mlir::createRecomposeONNXToONNXPass(
+        /*target=*/"", opts.enableRecomposeLayernormByTranspose));
+
   if (opts.enableONNXHybridPass) {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createONNXHybridTransformPass(
         !opts.disableRecomposeOption, opts.enableQuarkQuantizedLegalization,
         opts.enableConvTransposeDecompose,
         opts.enableConvTransposeDecomposeToPhasedConv,
-        opts.enableConvTranspose1dDecomposeToPhasedConv));
+        opts.enableConvTranspose1dDecomposeToPhasedConv, opts.enableRecomposeLayernormByTranspose));
     // Convolution Optimization for CPU: enable when there are no accelerators.
     if (targetCPU && opts.enableConvOptPass) {
       pm.addNestedPass<func::FuncOp>(onnx_mlir::createConvOptONNXToONNXPass(

src/Compiler/OnnxToMlirPasses.hpp

@@ -19,6 +19,7 @@ struct OnnxToMlirOptions {
   bool enableRemoveDqQOp = true;
   bool enableRemoveDqQAroundOp = true;
   bool enableRemoveBinary = false;
+  bool enableRecomposeLayernormByTranspose = false;
 
   bool disableRecomposeOption = false;
   bool enableONNXHybridPass = true;

src/Dialect/ONNX/DialectBuilder.cpp

@@ -591,6 +591,34 @@ Value OnnxBuilder::unsqueeze(Type outputType, Value data, Value axes) const {
       toTensor(outputType), toTensor(data), toTensor(axes));
 }
 
+Value OnnxBuilder::upRank(
+    mlir::Value data, int64_t toRank, bool trailing) const {
+  assert(data && "the value doesn't exist");
+
+  auto tensor = mlir::cast<mlir::ShapedType>(data.getType());
+  auto shape = getShape(tensor);
+  auto rank = getRank(tensor);
+  assert(rank <= toRank && "the rank of the tensor must be smaller");
+
+  if (rank == toRank)
+    return data;
+
+  int64_t rankDiff = toRank - rank;
+  SmallVector<int64_t> newShape;
+  if (trailing) {
+    newShape.append(shape.begin(), shape.end());
+    newShape.append(SmallVector<int64_t>(rankDiff, 1));
+  } else {
+    newShape.resize(rankDiff, 1);
+    newShape.append(shape.begin(), shape.end());
+  }
+
+  auto newType = tensor.clone(newShape);
+  auto shapeConst = constantInt64(newShape);
+  auto reshaped = reshape(newType, data, shapeConst);
+  return reshaped;
+}
+
 Value OnnxBuilder::where(
     Type outputType, Value condition, Value X, Value Y) const {
   return createTypedOpAndInferShapes<ONNXWhereOp>(

src/Dialect/ONNX/DialectBuilder.hpp

@@ -246,6 +246,11 @@ struct OnnxBuilder : DialectBuilder {
   mlir::Value unsqueeze(
       mlir::Type outputType, mlir::Value data, mlir::Value axes) const;
 
+  // Up ranking of the data tensor with reshape operator. The trailing is the
+  // option to choose to add the dimension with size 1 as leading or trailing.
+  mlir::Value upRank(
+      mlir::Value data, int64_t toRank, bool trailing = false) const;
+
   // ONNXWhereOp
   mlir::Value where(mlir::Type outputType, mlir::Value condition, mlir::Value X,
       mlir::Value Y) const;

src/Dialect/ONNX/ONNXOps/Canonicalize.cpp

@@ -20,6 +20,7 @@
 #include <numeric>
 
 #include "mlir/Dialect/Traits.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
@@ -584,6 +585,66 @@ class PropagateReshapeThroughBinaryOpPattern
   };
 };
 
+// This pattern bubbles up AddOp through transpose to keep the bias Add
+// operation right after LN_type op. This will helps the other patterns fold the
+// add into the operands of LayerNorm.
+//
+// From: LayerNorm
+//    |
+// Transpose
+//    |
+//   Add
+//
+// To:
+// LayerNorm
+//    |
+//   Add
+//    |
+// Transpose
+template <typename LN_TYPE>
+class BubbleUpBiasForLayerNormPattern : public OpRewritePattern<ONNXAddOp> {
+public:
+  using OpRewritePattern<ONNXAddOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(
+      ONNXAddOp addOp, PatternRewriter &r) const override {
+    if (!isConstLikeValue(addOp.getB()))
+      return r.notifyMatchFailure(addOp, "not a constant rhs operand");
+
+    auto transposeOp =
+        llvm::dyn_cast_or_null<ONNXTransposeOp>(addOp.getA().getDefiningOp());
+    if (!transposeOp)
+      return r.notifyMatchFailure(addOp, "the producer is not a transpose");
+
+    if (!transposeOp->hasOneUse())
+      return r.notifyMatchFailure(
+          addOp, "cannot bubble up because transpose has other user");
+
+    auto layernormResult = transposeOp.getData();
+    auto layerNorm =
+        llvm::dyn_cast_or_null<LN_TYPE>(layernormResult.getDefiningOp());
+    if (!layerNorm)
+      return r.notifyMatchFailure(
+          transposeOp, "the producer is not a layernorm");
+
+    if (!isNoneValue(layerNorm.getB()))
+      return r.notifyMatchFailure(layerNorm, "layernorm already has a bias");
+
+    OnnxBuilder create(r, addOp.getLoc());
+
+    auto perm = extractFromIntegerArrayAttr<int64_t>(transposeOp.getPermAttr());
+    auto invertedPerm = invertPermutationVector(perm);
+    auto cstReshaped = create.upRank(addOp.getB(), getRank(addOp.getType()));
+    auto cstTranposed = create.transposeInt64(cstReshaped, invertedPerm);
+    auto newAddOp = create.add(layernormResult, cstTranposed);
+    auto transposedBack = create.transposeInt64(newAddOp, perm);
+
+    r.replaceOp(addOp, transposedBack);
+
+    return success();
+  };
+};
+
 // This rewriting is to optimize the scalar Div/Mul in self-attention layers.
 // In particular, it rewrites the following pattern:
 // ```
@@ -2426,6 +2487,10 @@ void ONNXAddOp::getCanonicalizationPatterns(
       PropagateBiasIntoLayerNormRewritePattern<ONNXRMSLayerNormalizationOp>>(
       context);
   results.insert<PropagateReshapeThroughBinaryOpPattern<ONNXAddOp>>(context);
+  results.insert<BubbleUpBiasForLayerNormPattern<ONNXLayerNormalizationOp>>(
+      context);
+  results.insert<BubbleUpBiasForLayerNormPattern<ONNXRMSLayerNormalizationOp>>(
+      context);
 }
 
 /// on the ONNXAndOp.

src/Dialect/ONNX/Transforms/ONNXHybridTransformPass.cpp

@@ -111,20 +111,28 @@ struct ONNXHybridTransformPass
           "phased Conv"),
       ::llvm::cl::init(false)};
 
+  Option<bool> recomposeLayernormByTranspose{*this,
+      "recompose-layernorm-by-transpose",
+      llvm::cl::desc("Use transpose operator to make unsuitable axes suitable "
+                     "for matching layernorm"),
+      ::llvm::cl::init(false)};
+
   FrozenRewritePatternSet patterns;
 
   ONNXHybridTransformPass(bool enableRecomposition,
       bool enableQuarkQuantizedOpsLegalization,
       bool enableConvTransposeDecompose,
       bool enableConvTransposeDecomposeToPhasedConv,
-      bool enableConvTranspose1dDecomposeToPhasedConv) {
+      bool enableConvTranspose1dDecomposeToPhasedConv,
+      bool recomposeLayernormByTranspose) {
     this->recomposition = enableRecomposition;
     this->quarkQuantizedOpsLegalization = enableQuarkQuantizedOpsLegalization;
     this->enableConvTransposeDecompose = enableConvTransposeDecompose;
     this->enableConvTransposeDecomposeToPhasedConv =
         enableConvTransposeDecomposeToPhasedConv;
     this->enableConvTranspose1dDecomposeToPhasedConv =
         enableConvTranspose1dDecomposeToPhasedConv;
+    this->recomposeLayernormByTranspose = recomposeLayernormByTranspose;
   }
 
   ONNXHybridTransformPass(const ONNXHybridTransformPass &pass)
@@ -171,7 +179,8 @@ struct ONNXHybridTransformPass
     }
 
     if (recomposition) {
-      getRecomposeONNXToONNXPatterns(cumulativePatterns);
+      getRecomposeONNXToONNXPatterns(
+          cumulativePatterns, recomposeLayernormByTranspose);
     }
 
     patterns = FrozenRewritePatternSet(std::move(cumulativePatterns));
@@ -210,9 +219,11 @@ std::unique_ptr<mlir::Pass> onnx_mlir::createONNXHybridTransformPass(
     bool enableRecomposition, bool enableQuarkQuantizedOpsLegalization,
     bool enableConvTransposeDecompose,
     bool enableConvTransposeDecomposeToPhasedConv,
-    bool enableConvTranspose1dDecomposeToPhasedConv) {
+    bool enableConvTranspose1dDecomposeToPhasedConv,
+    bool enableRecomposeLayernormByTranspose) {
   return std::make_unique<ONNXHybridTransformPass>(enableRecomposition,
       enableQuarkQuantizedOpsLegalization, enableConvTransposeDecompose,
       enableConvTransposeDecomposeToPhasedConv,
-      enableConvTranspose1dDecomposeToPhasedConv);
+      enableConvTranspose1dDecomposeToPhasedConv,
+      enableRecomposeLayernormByTranspose);
 }