Merge pull request #465 from Xilinx/move_instancenorm_pattern

Refactor: Move InstanceNorm→LayerNorm from Canonicalize to Decompose

Author: ayanshul-amd

src/Compiler/OnnxToMlirPasses.cpp

@@ -30,7 +30,8 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
   pm.addNestedPass<func::FuncOp>(onnx_mlir::createDecomposeONNXToONNXPass(
       /*target=*/"", opts.enableConvTransposeDecompose,
       opts.enableConvTransposeDecomposeToPhasedConv,
-      opts.enableConvTranspose1dDecomposeToPhasedConv));
+      opts.enableConvTranspose1dDecomposeToPhasedConv,
+      opts.enableInstanceNormDecompose));
   if (!opts.disableRecomposeOption)
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createRecomposeONNXToONNXPass(
         /*target=*/"", opts.enableRecomposeLayernormByTranspose));
@@ -41,7 +42,8 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
         opts.enableConvTransposeDecompose,
         opts.enableConvTransposeDecomposeToPhasedConv,
         opts.enableConvTranspose1dDecomposeToPhasedConv,
-        opts.enableRecomposeLayernormByTranspose));
+        opts.enableRecomposeLayernormByTranspose,
+        opts.enableInstanceNormDecompose));
     // Convolution Optimization for CPU: enable when there are no accelerators.
     if (targetCPU && opts.enableConvOptPass) {
       pm.addNestedPass<func::FuncOp>(onnx_mlir::createConvOptONNXToONNXPass(
@@ -52,7 +54,8 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
               opts.enableConvTransposeDecompose,
               opts.enableConvTransposeDecomposeToPhasedConv,
               opts.enableConvTranspose1dDecomposeToPhasedConv,
-              opts.enableRecomposeLayernormByTranspose));
+              opts.enableRecomposeLayernormByTranspose,
+              opts.enableInstanceNormDecompose));
     }
   } else {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
@@ -109,7 +112,8 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
         opts.enableConvTransposeDecompose,
         opts.enableConvTransposeDecomposeToPhasedConv,
         opts.enableConvTranspose1dDecomposeToPhasedConv,
-        opts.enableRecomposeLayernormByTranspose));
+        opts.enableRecomposeLayernormByTranspose,
+        opts.enableInstanceNormDecompose));
   } else {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
     pm.addPass(mlir::createCanonicalizerPass());

src/Compiler/OnnxToMlirPasses.hpp

@@ -16,6 +16,7 @@ struct OnnxToMlirOptions {
   bool enableConvTransposeDecompose = false;
   bool enableConvTransposeDecomposeToPhasedConv = false;
   bool enableConvTranspose1dDecomposeToPhasedConv = false;
+  bool enableInstanceNormDecompose = true;
   bool enableRemoveDqQOp = true;
   bool enableRemoveDqQAroundOp = true;
   bool enableRemoveBinary = false;

src/Dialect/ONNX/ONNXOps.td.inc

@@ -3643,7 +3643,6 @@ def ONNXIfOp:ONNX_Op<"If",
 
 def ONNXInstanceNormalizationOp:ONNX_Op<"InstanceNormalization",
   [Pure, OpVersionTrait<22>, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>, DeclareOpInterfaceMethods<ShapeHelperOpInterface>]> {
-  let hasCanonicalizer = 1;
   let summary = "ONNX InstanceNormalization operation";
   let description = [{
   Carries out instance normalization as described in the paper

src/Dialect/ONNX/ONNXOps/Canonicalize.cpp

@@ -2087,68 +2087,6 @@ class RemoveWhereEqualPattern : public OpRewritePattern<ONNXWhereOp> {
   }
 };
 
-// =============================================================================
-// Rewrite pattern for Instance Normalization
-// =============================================================================
-
-struct RemoveInstanceNormPattern
-    : public OpRewritePattern<ONNXInstanceNormalizationOp> {
-  using OpRewritePattern<ONNXInstanceNormalizationOp>::OpRewritePattern;
-
-  static bool isDecomposable(ONNXInstanceNormalizationOp instanceNormOp) {
-    return onnx_mlir::hasStaticShape(instanceNormOp.getInput().getType()) &&
-           onnx_mlir::hasStaticShape(instanceNormOp.getOutput().getType());
-  }
-
-  LogicalResult matchAndRewrite(ONNXInstanceNormalizationOp instanceNormOp,
-      PatternRewriter &rewriter) const final {
-    // Match.
-    if (!isDecomposable(instanceNormOp)) {
-      return failure();
-    }
-
-    // Get info.
-    Value input = instanceNormOp.getInput();
-    Value scale = instanceNormOp.getScale();
-    Value bias = instanceNormOp.getB();
-    ShapedType inputType = mlir::cast<ShapedType>(input.getType());
-    Type elementType = inputType.getElementType();
-    auto inputShape = inputType.getShape();
-    int64_t C = inputShape[1];
-    int64_t inputRank = inputType.getRank();
-    int64_t nonSpacialRank = 2; //  Batch N and Channel C: 2 dimensions.
-    assert(inputRank > nonSpacialRank &&
-           "expected instance norm with input ranks > 2");
-
-    // Rewrite.
-    onnx_mlir::MultiDialectBuilder<onnx_mlir::OnnxBuilder> create(
-        rewriter, instanceNormOp.getLoc());
-    int64_t axis = nonSpacialRank;
-    int64_t numInNorm = inputRank - axis;
-    // Unsqueeze scale/bias from [C] to [C x 1 x 1 x ... x 1] with numInNorm
-    // 1s.
-    llvm::SmallVector<int64_t, 4> axesList, biasScaleShape;
-    biasScaleShape.emplace_back(C);
-    for (int64_t i = 1; i <= numInNorm; ++i) {
-      biasScaleShape.emplace_back(1);
-      axesList.emplace_back(i);
-    }
-    Value axes = create.onnx.constantInt64(axesList);
-    Type biasScaleType = RankedTensorType::get(biasScaleShape, elementType);
-    Value newScale = create.onnx.unsqueeze(biasScaleType, scale, axes);
-    Value newBias = create.onnx.unsqueeze(biasScaleType, bias, axes);
-    // Create output using layer norm.
-    Value Y = create.onnx.layerNorm(inputType, input, newScale, newBias, axis,
-        instanceNormOp.getEpsilonAttr());
-    // Set the type of the output to be the same as the output of the original
-    // operation we are trying to replace.
-    Y.setType(instanceNormOp.getResult().getType());
-    // Replace operation.
-    rewriter.replaceOp(instanceNormOp, Y);
-    return success();
-  }
-};
-
 // =============================================================================
 // Rewrite pattern for Group Normalization
 // =============================================================================
@@ -2614,12 +2552,6 @@ void ONNXIdentityOp::getCanonicalizationPatterns(
   results.insert<IdentityEliminationPattern>(context);
 }
 
-/// on the ONNXInstanceNormalizationOp.
-void ONNXInstanceNormalizationOp::getCanonicalizationPatterns(
-    RewritePatternSet &results, MLIRContext *context) {
-  results.insert<RemoveInstanceNormPattern>(context);
-}
-
 /// on the ONNXLayoutTransformOp.
 void ONNXLayoutTransformOp::getCanonicalizationPatterns(
     RewritePatternSet &result, MLIRContext *context) {

src/Dialect/ONNX/Transforms/Decompose.cpp

@@ -3512,6 +3512,67 @@ class ReplaceCastLikeByCastPattern : public OpRewritePattern<ONNXCastLikeOp> {
   }
 };
 
+// =============================================================================
+// Decompose InstanceNormalization to LayerNormalization
+// =============================================================================
+struct DecomposeInstanceNormPattern
+    : public OpRewritePattern<ONNXInstanceNormalizationOp> {
+  using OpRewritePattern<ONNXInstanceNormalizationOp>::OpRewritePattern;
+
+  static bool isDecomposable(ONNXInstanceNormalizationOp instanceNormOp) {
+    return onnx_mlir::hasStaticShape(instanceNormOp.getInput().getType()) &&
+           onnx_mlir::hasStaticShape(instanceNormOp.getOutput().getType());
+  }
+
+  LogicalResult matchAndRewrite(ONNXInstanceNormalizationOp instanceNormOp,
+      PatternRewriter &rewriter) const final {
+    // Match.
+    if (!isDecomposable(instanceNormOp)) {
+      return failure();
+    }
+
+    // Get info.
+    Value input = instanceNormOp.getInput();
+    Value scale = instanceNormOp.getScale();
+    Value bias = instanceNormOp.getB();
+    ShapedType inputType = mlir::cast<ShapedType>(input.getType());
+    Type elementType = inputType.getElementType();
+    auto inputShape = inputType.getShape();
+    int64_t C = inputShape[1];
+    int64_t inputRank = inputType.getRank();
+    int64_t nonSpacialRank = 2; //  Batch N and Channel C: 2 dimensions.
+    assert(inputRank > nonSpacialRank &&
+           "expected instance norm with input ranks > 2");
+
+    // Rewrite.
+    onnx_mlir::MultiDialectBuilder<onnx_mlir::OnnxBuilder> create(
+        rewriter, instanceNormOp.getLoc());
+    int64_t axis = nonSpacialRank;
+    int64_t numInNorm = inputRank - axis;
+    // Unsqueeze scale/bias from [C] to [C x 1 x 1 x ... x 1] with numInNorm
+    // 1s.
+    llvm::SmallVector<int64_t, 4> axesList, biasScaleShape;
+    biasScaleShape.emplace_back(C);
+    for (int64_t i = 1; i <= numInNorm; ++i) {
+      biasScaleShape.emplace_back(1);
+      axesList.emplace_back(i);
+    }
+    Value axes = create.onnx.constantInt64(axesList);
+    Type biasScaleType = RankedTensorType::get(biasScaleShape, elementType);
+    Value newScale = create.onnx.unsqueeze(biasScaleType, scale, axes);
+    Value newBias = create.onnx.unsqueeze(biasScaleType, bias, axes);
+    // Create output using layer norm.
+    Value Y = create.onnx.layerNorm(inputType, input, newScale, newBias, axis,
+        instanceNormOp.getEpsilonAttr());
+    // Set the type of the output to be the same as the output of the original
+    // operation we are trying to replace.
+    Y.setType(instanceNormOp.getResult().getType());
+    // Replace operation.
+    rewriter.replaceOp(instanceNormOp, Y);
+    return success();
+  }
+};
+
 // =============================================================================
 // Decompose Hardswish to simpler ONNX ops
 // =============================================================================
@@ -3577,13 +3638,15 @@ struct DecomposeONNXToONNXPass
   DecomposeONNXToONNXPass(const std::string &target,
       bool enableConvTransposeDecompose = false,
       bool enableConvTransposeDecomposeToPhasedConv = false,
-      bool enableConvTranspose1dDecomposeToPhasedConv = false) {
+      bool enableConvTranspose1dDecomposeToPhasedConv = false,
+      bool enableInstanceNormDecompose = true) {
     this->target = target;
     this->enableConvTransposeDecompose = enableConvTransposeDecompose;
     this->enableConvTransposeDecomposeToPhasedConv =
         enableConvTransposeDecomposeToPhasedConv;
     this->enableConvTranspose1dDecomposeToPhasedConv =
         enableConvTranspose1dDecomposeToPhasedConv;
+    this->enableInstanceNormDecompose = enableInstanceNormDecompose;
   }
 
   DecomposeONNXToONNXPass(const DecomposeONNXToONNXPass &pass)
@@ -3594,6 +3657,8 @@ struct DecomposeONNXToONNXPass
         pass.enableConvTransposeDecompose.getValue();
     this->enableConvTransposeDecomposeToPhasedConv =
         pass.enableConvTransposeDecomposeToPhasedConv.getValue();
+    this->enableInstanceNormDecompose =
+        pass.enableInstanceNormDecompose.getValue();
   }
 
   StringRef getArgument() const override { return "decompose-onnx"; }
@@ -3623,6 +3688,12 @@ struct DecomposeONNXToONNXPass
           "phased Conv"),
       ::llvm::cl::init(false)};
 
+  Option<bool> enableInstanceNormDecompose{*this,
+      "enable-instancenorm-decompose",
+      llvm::cl::desc("Enable decomposition of InstanceNormalization to "
+                     "LayerNormalization"),
+      ::llvm::cl::init(true)};
+
   void runOnOperation() final;
 
   typedef PassWrapper<DecomposeONNXToONNXPass, OperationPass<func::FuncOp>>
@@ -3635,7 +3706,7 @@ void DecomposeONNXToONNXPass::runOnOperation() {
   RewritePatternSet patterns(context);
   onnx_mlir::getDecomposeONNXToONNXPatterns(patterns,
       enableConvTransposeDecompose, enableConvTransposeDecomposeToPhasedConv,
-      enableConvTranspose1dDecomposeToPhasedConv);
+      enableConvTranspose1dDecomposeToPhasedConv, enableInstanceNormDecompose);
   patterns.insert<ReplaceCastLikeByCastPattern>(context);
 
 #ifdef ONNX_MLIR_ENABLE_STABLEHLO
@@ -3653,7 +3724,8 @@ void DecomposeONNXToONNXPass::runOnOperation() {
 void onnx_mlir::getDecomposeONNXToONNXPatterns(
     mlir::RewritePatternSet &patterns, bool enableConvTransposeDecompose,
     bool enableConvTransposeDecomposeToPhasedConv,
-    bool enableConvTranspose1dDecomposeToPhasedConv) {
+    bool enableConvTranspose1dDecomposeToPhasedConv,
+    bool enableInstanceNormDecompose) {
   MLIRContext *context = patterns.getContext();
   populateWithGenerated(patterns);
   if (enableConvTransposeDecompose)
@@ -3662,6 +3734,8 @@ void onnx_mlir::getDecomposeONNXToONNXPatterns(
     convtranspose_phased::populateWithGenerated(patterns);
   if (enableConvTranspose1dDecomposeToPhasedConv)
     convtranspose_1d_phased::populateWithGenerated(patterns);
+  if (enableInstanceNormDecompose)
+    patterns.insert<DecomposeInstanceNormPattern>(context);
   patterns.insert<onnx_mlir::DecomposeEinsumPattern>(context);
   patterns.insert<ConcatFusePattern>(context);
   patterns.insert<DecomposeHardSwishPattern>(context);
@@ -3699,8 +3773,9 @@ void onnx_mlir::getDecomposeONNXToONNXPatterns(
 std::unique_ptr<mlir::Pass> onnx_mlir::createDecomposeONNXToONNXPass(
     const std::string &target, bool enableConvTransposeDecompose,
     bool enableConvTransposeDecomposeToPhasedConv,
-    bool enableConvTranspose1dDecomposeToPhasedConv) {
+    bool enableConvTranspose1dDecomposeToPhasedConv,
+    bool enableInstanceNormDecompose) {
   return std::make_unique<DecomposeONNXToONNXPass>(target,
       enableConvTransposeDecompose, enableConvTransposeDecomposeToPhasedConv,
-      enableConvTranspose1dDecomposeToPhasedConv);
+      enableConvTranspose1dDecomposeToPhasedConv, enableInstanceNormDecompose);
 }

src/Dialect/ONNX/Transforms/Decompose.hpp

@@ -29,7 +29,8 @@ namespace onnx_mlir {
 void getDecomposeONNXToONNXPatterns(mlir::RewritePatternSet &patterns,
     bool enableConvTransposeDecompose,
     bool enableConvTransposeDecomposeToPhasedConv,
-    bool enableConvTranspose1dDecomposeToPhasedConv);
+    bool enableConvTranspose1dDecomposeToPhasedConv,
+    bool enableInstanceNormDecompose);
 
 } // namespace onnx_mlir
 #endif

src/Dialect/ONNX/Transforms/ONNXHybridTransformPass.cpp

@@ -111,6 +111,12 @@ struct ONNXHybridTransformPass
           "phased Conv"),
       ::llvm::cl::init(false)};
 
+  Option<bool> enableInstanceNormDecompose{*this,
+      "enable-instancenorm-decompose",
+      llvm::cl::desc("Enable decomposition of InstanceNormalization to "
+                     "LayerNormalization"),
+      ::llvm::cl::init(true)};
+
   Option<bool> recomposeLayernormByTranspose{*this,
       "recompose-layernorm-by-transpose",
       llvm::cl::desc("Use transpose operator to make unsuitable axes suitable "
@@ -124,7 +130,7 @@ struct ONNXHybridTransformPass
       bool enableConvTransposeDecompose,
       bool enableConvTransposeDecomposeToPhasedConv,
       bool enableConvTranspose1dDecomposeToPhasedConv,
-      bool recomposeLayernormByTranspose) {
+      bool recomposeLayernormByTranspose, bool enableInstanceNormDecompose) {
     this->recomposition = enableRecomposition;
     this->quarkQuantizedOpsLegalization = enableQuarkQuantizedOpsLegalization;
     this->enableConvTransposeDecompose = enableConvTransposeDecompose;
@@ -133,6 +139,7 @@ struct ONNXHybridTransformPass
     this->enableConvTranspose1dDecomposeToPhasedConv =
         enableConvTranspose1dDecomposeToPhasedConv;
     this->recomposeLayernormByTranspose = recomposeLayernormByTranspose;
+    this->enableInstanceNormDecompose = enableInstanceNormDecompose;
   }
 
   ONNXHybridTransformPass(const ONNXHybridTransformPass &pass)
@@ -175,7 +182,8 @@ struct ONNXHybridTransformPass
       getDecomposeONNXToONNXPatterns(cumulativePatterns,
           enableConvTransposeDecompose,
           enableConvTransposeDecomposeToPhasedConv,
-          enableConvTranspose1dDecomposeToPhasedConv);
+          enableConvTranspose1dDecomposeToPhasedConv,
+          enableInstanceNormDecompose);
     }
 
     if (recomposition) {
@@ -220,10 +228,11 @@ std::unique_ptr<mlir::Pass> onnx_mlir::createONNXHybridTransformPass(
     bool enableConvTransposeDecompose,
     bool enableConvTransposeDecomposeToPhasedConv,
     bool enableConvTranspose1dDecomposeToPhasedConv,
-    bool enableRecomposeLayernormByTranspose) {
+    bool enableRecomposeLayernormByTranspose,
+    bool enableInstanceNormDecompose) {
   return std::make_unique<ONNXHybridTransformPass>(enableRecomposition,
       enableQuarkQuantizedOpsLegalization, enableConvTransposeDecompose,
       enableConvTransposeDecomposeToPhasedConv,
       enableConvTranspose1dDecomposeToPhasedConv,
-      enableRecomposeLayernormByTranspose);
+      enableRecomposeLayernormByTranspose, enableInstanceNormDecompose);
 }

src/Pass/Passes.hpp

@@ -40,7 +40,8 @@ std::unique_ptr<mlir::Pass> createONNXOpTransformPass(int threshold,
 std::unique_ptr<mlir::Pass> createDecomposeONNXToONNXPass(
     const std::string &target = "", bool enableConvTransposeDecompose = false,
     bool enableConvTransposeDecomposeToPhasedConv = false,
-    bool enableConvTranspose1dDecomposeToPhasedConv = false);
+    bool enableConvTranspose1dDecomposeToPhasedConv = false,
+    bool enableInstanceNormDecompose = true);
 std::unique_ptr<mlir::Pass> createRecomposeONNXToONNXPass(
     const std::string &target = "",
     const bool &recomposeLayernormByTranspose = false);
@@ -86,7 +87,8 @@ std::unique_ptr<mlir::Pass> createONNXHybridTransformPass(
     bool enableConvTransposeDecompose = false,
     bool enableConvTransposeDecomposeToPhasedConv = false,
     bool enableConvTranspose1dDecomposeToPhasedConv = false,
-    bool enableRecomposeLayernormByTranspose = false);
+    bool enableRecomposeLayernormByTranspose = false,
+    bool enableInstanceNormDecompose = true);
 
 /// Pass for analyzing unknown dimension in ONNX operations.
 std::unique_ptr<mlir::Pass> createONNXDimAnalysisPass();

test/mlir/onnx/onnx_canonicalization.mlir

@@ -2133,22 +2133,6 @@ func.func @test_reorder_relu_maxpool(%arg0: tensor<1x64x32x32xf32>) -> tensor<1x
 
 // -----
 
-func.func @test_instancenorm(%arg0: tensor<2x3x4x5x6xf32>, %arg1: tensor<3xf32>, %arg2: tensor<3xf32>) -> tensor<2x3x4x5x6xf32> {
-  %0 = "onnx.InstanceNormalization"(%arg0, %arg1, %arg2) {epsilon = 0.00999999977 : f32} : (tensor<2x3x4x5x6xf32>, tensor<3xf32>, tensor<3xf32>) -> tensor<2x3x4x5x6xf32>
-  onnx.Return %0 : tensor<2x3x4x5x6xf32>
-// mlir2FileCheck.py
-// CHECK-LABEL:  func.func @test_instancenorm
-// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<2x3x4x5x6xf32>, [[PARAM_1_:%.+]]: tensor<3xf32>, [[PARAM_2_:%.+]]: tensor<3xf32>) -> tensor<2x3x4x5x6xf32> {
-// CHECK:           [[VAR_0_:%.+]] = onnx.Constant dense<[1, 2, 3]> : tensor<3xi64>
-// CHECK-DAG:       [[VAR_1_:%.+]] = "onnx.Unsqueeze"([[PARAM_1_]], [[VAR_0_]]) : (tensor<3xf32>, tensor<3xi64>) -> tensor<3x1x1x1xf32>
-// CHECK-DAG:       [[VAR_2_:%.+]] = "onnx.Unsqueeze"([[PARAM_2_]], [[VAR_0_]]) : (tensor<3xf32>, tensor<3xi64>) -> tensor<3x1x1x1xf32>
-// CHECK:           [[Y_:%.+]], [[Mean_:%.+]], [[VAR_InvStdDev_:%.+]] = "onnx.LayerNormalization"([[PARAM_0_]], [[VAR_1_]], [[VAR_2_]]) {axis = 2 : si64, epsilon = 0.00999999977 : f32, stash_type = 1 : si64} : (tensor<2x3x4x5x6xf32>, tensor<3x1x1x1xf32>, tensor<3x1x1x1xf32>) -> (tensor<2x3x4x5x6xf32>, none, none)
-// CHECK:           onnx.Return [[Y_]] : tensor<2x3x4x5x6xf32>
-// CHECK:         }
-}
-
-// -----
-
 func.func @test_groupnorm_v18(%arg0: tensor<3x4x2x2xf32>, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>) -> tensor<3x4x2x2xf32> {
   %0 = "onnx.GroupNormalizationV18"(%arg0, %arg1, %arg2) {epsilon = 0.00999999977 : f32, num_groups = 2 : si64} : (tensor<3x4x2x2xf32>, tensor<2xf32>, tensor<2xf32>) -> tensor<3x4x2x2xf32>
   onnx.Return %0 : tensor<3x4x2x2xf32>