Address review comments: move enableInstanceNormDecompose to end and add e2e test

Author: ayanshul-amd

src/Compiler/OnnxToMlirPasses.cpp

@@ -42,8 +42,8 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
         opts.enableConvTransposeDecompose,
         opts.enableConvTransposeDecomposeToPhasedConv,
         opts.enableConvTranspose1dDecomposeToPhasedConv,
-        opts.enableInstanceNormDecompose,
-        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(
@@ -54,8 +54,8 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
               opts.enableConvTransposeDecompose,
               opts.enableConvTransposeDecomposeToPhasedConv,
               opts.enableConvTranspose1dDecomposeToPhasedConv,
-              opts.enableInstanceNormDecompose,
-              opts.enableRecomposeLayernormByTranspose));
+              opts.enableRecomposeLayernormByTranspose,
+              opts.enableInstanceNormDecompose));
     }
   } else {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
@@ -112,8 +112,8 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
         opts.enableConvTransposeDecompose,
         opts.enableConvTransposeDecomposeToPhasedConv,
         opts.enableConvTranspose1dDecomposeToPhasedConv,
-        opts.enableInstanceNormDecompose,
-        opts.enableRecomposeLayernormByTranspose));
+        opts.enableRecomposeLayernormByTranspose,
+        opts.enableInstanceNormDecompose));
   } else {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
     pm.addPass(mlir::createCanonicalizerPass());

src/Dialect/ONNX/Transforms/ONNXHybridTransformPass.cpp

@@ -130,16 +130,16 @@ struct ONNXHybridTransformPass
       bool enableConvTransposeDecompose,
       bool enableConvTransposeDecomposeToPhasedConv,
       bool enableConvTranspose1dDecomposeToPhasedConv,
-      bool enableInstanceNormDecompose, bool recomposeLayernormByTranspose) {
+      bool recomposeLayernormByTranspose, bool enableInstanceNormDecompose) {
     this->recomposition = enableRecomposition;
     this->quarkQuantizedOpsLegalization = enableQuarkQuantizedOpsLegalization;
     this->enableConvTransposeDecompose = enableConvTransposeDecompose;
     this->enableConvTransposeDecomposeToPhasedConv =
         enableConvTransposeDecomposeToPhasedConv;
     this->enableConvTranspose1dDecomposeToPhasedConv =
         enableConvTranspose1dDecomposeToPhasedConv;
-    this->enableInstanceNormDecompose = enableInstanceNormDecompose;
     this->recomposeLayernormByTranspose = recomposeLayernormByTranspose;
+    this->enableInstanceNormDecompose = enableInstanceNormDecompose;
   }
 
   ONNXHybridTransformPass(const ONNXHybridTransformPass &pass)
@@ -228,11 +228,11 @@ std::unique_ptr<mlir::Pass> onnx_mlir::createONNXHybridTransformPass(
     bool enableConvTransposeDecompose,
     bool enableConvTransposeDecomposeToPhasedConv,
     bool enableConvTranspose1dDecomposeToPhasedConv,
-    bool enableInstanceNormDecompose,
-    bool enableRecomposeLayernormByTranspose) {
+    bool enableRecomposeLayernormByTranspose,
+    bool enableInstanceNormDecompose) {
   return std::make_unique<ONNXHybridTransformPass>(enableRecomposition,
       enableQuarkQuantizedOpsLegalization, enableConvTransposeDecompose,
       enableConvTransposeDecomposeToPhasedConv,
-      enableConvTranspose1dDecomposeToPhasedConv, enableInstanceNormDecompose,
-      enableRecomposeLayernormByTranspose);
+      enableConvTranspose1dDecomposeToPhasedConv,
+      enableRecomposeLayernormByTranspose, enableInstanceNormDecompose);
 }

src/Pass/Passes.hpp

@@ -87,8 +87,8 @@ std::unique_ptr<mlir::Pass> createONNXHybridTransformPass(
     bool enableConvTransposeDecompose = false,
     bool enableConvTransposeDecomposeToPhasedConv = false,
     bool enableConvTranspose1dDecomposeToPhasedConv = false,
-    bool enableInstanceNormDecompose = true,
-    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/parse/instancenorm_to_layernorm.onnxtext

@@ -0,0 +1,19 @@
+// RUN: onnx-mlir --EmitONNXBasic --printIR %s | FileCheck %s
+<
+   ir_version: 8,
+   opset_import: ["" : 6]
+>
+test_instancenorm_e2e (float[2,3,4,5] input, float[3] scale, float[3] bias) => (float[2,3,4,5] output) {
+   output = InstanceNormalization <epsilon: float = 0.01> (input, scale, bias)
+}
+
+// mlir2FileCheck.py
+// CHECK-LABEL:  func.func @main_graph
+// CHECK-SAME:   ([[INPUT:%.+]]: tensor<2x3x4x5xf32> {onnx.name = "input"}, [[SCALE:%.+]]: tensor<3xf32> {onnx.name = "scale"}, [[BIAS:%.+]]: tensor<3xf32> {onnx.name = "bias"}) -> (tensor<2x3x4x5xf32> {onnx.name = "output"}) {
+// CHECK-DAG:       [[AXES:%.+]] = onnx.Constant dense<[1, 2]> : tensor<2xi64>
+// CHECK-DAG:       [[UNSQUEEZE_SCALE:%.+]] = "onnx.Unsqueeze"([[SCALE]], [[AXES]]) : (tensor<3xf32>, tensor<2xi64>) -> tensor<3x1x1xf32>
+// CHECK-DAG:       [[UNSQUEEZE_BIAS:%.+]] = "onnx.Unsqueeze"([[BIAS]], [[AXES]]) : (tensor<3xf32>, tensor<2xi64>) -> tensor<3x1x1xf32>
+// CHECK:           [[OUTPUT:%.+]], [[MEAN:%.+]], [[INV_STD_DEV:%.+]] = "onnx.LayerNormalization"([[INPUT]], [[UNSQUEEZE_SCALE]], [[UNSQUEEZE_BIAS]]) {axis = 2 : si64, epsilon = 1.000000e-02 : f32, stash_type = 1 : si64} : (tensor<2x3x4x5xf32>, tensor<3x1x1xf32>, tensor<3x1x1xf32>) -> (tensor<2x3x4x5xf32>, none, none)
+// CHECK:           onnx.Return [[OUTPUT]] : tensor<2x3x4x5xf32>
+// CHECK:         }
+