Layernorm: convert instance norm and group norm to layer norm. #2595

AlexandreEichenberger · 2023-10-30T15:28:20Z

Layer norm is a superset of previous InstanceNormalization and GroupNormalization.

For instance norm:

    %0 = "onnx.InstanceNormalization"(%arg0, %arg1, %arg2) {epsilon = 0.00999999977 : f32} : (tensor<2x3x4x5x6xf32>, tensor<3xf32>, tensor<3xf32>) -> tensor<2x3x4x5x6xf32>

becomes

  %0 = onnx.Constant dense<[1, 2, 3]> : tensor<3xi64>
  %1 = "onnx.Unsqueeze"(%arg1, %0) : (tensor<3xf32>, tensor<3xi64>) -> tensor<3x1x1x1xf32>
  %2 = "onnx.Unsqueeze"(%arg2, %0) : (tensor<3xf32>, tensor<3xi64>) -> tensor<3x1x1x1xf32>
  %Y, %Mean, %InvStdDev = "onnx.LayerNormalization"(%arg0, %1, %2) {axis = 2 : si64, epsilon = 0.00999999977 : f32, stash_type = 1 : si64} : (tensor<2x3x4x5x6xf32>, tensor<3x1x1x1xf32>, tensor<3x1x1x1xf32>) -> (tensor<2x3x4x5x6xf32>, none, none)

For group norm:

    %0 = "onnx.GroupNormalization"(%arg0, %arg1, %arg2) {epsilon = 0.00999999977 : f32, num_groups = 2 : si64} : (tensor<3x4x6x8x16xf32>, tensor<2xf32>, tensor<2xf32>) -> tensor<3x4x6x8x16xf32>

becomes

  %0 = onnx.Constant dense<[3, 4, 6, 8, 16]> : tensor<5xi64>
  %1 = onnx.Constant dense<[3, 2, -1, 6, 8, 16]> : tensor<6xi64>
  %2 = onnx.Constant dense<[1, 2, 3, 4]> : tensor<4xi64>
  %3 = "onnx.Unsqueeze"(%arg1, %2) : (tensor<2xf32>, tensor<4xi64>) -> tensor<2x1x1x1x1xf32>
  %4 = "onnx.Unsqueeze"(%arg2, %2) : (tensor<2xf32>, tensor<4xi64>) -> tensor<2x1x1x1x1xf32>
  %5 = "onnx.Reshape"(%arg0, %1) {allowzero = 0 : si64} : (tensor<3x4x6x8x16xf32>, tensor<6xi64>) -> tensor<3x2x2x6x8x16xf32>
  %Y, %Mean, %InvStdDev = "onnx.LayerNormalization"(%5, %3, %4) {axis = 2 : si64, epsilon = 0.00999999977 : f32, stash_type = 1 : si64} : (tensor<3x2x2x6x8x16xf32>, tensor<2x1x1x1x1xf32>, tensor<2x1x1x1x1xf32>) -> (tensor<3x2x2x6x8x16xf32>, none, none)
  %6 = "onnx.Reshape"(%Y, %0) {allowzero = 0 : si64} : (tensor<3x2x2x6x8x16xf32>, tensor<5xi64>) -> tensor<3x4x6x8x16xf32>

Signed-off-by: Alexandre Eichenberger <[email protected]>

…onnx-mlir into layernorm

Signed-off-by: Alexandre Eichenberger <[email protected]>

…er/onnx-mlir into layernorm_v2

Signed-off-by: Alexandre Eichenberger <[email protected]>

AlexandreEichenberger · 2023-10-30T15:29:20Z

@philass @sorenlassen : this transformation is done at decompose, and at this time, all instance and group norms are switched to layer norm.

Let me know if you want a switch.

Signed-off-by: Alexandre Eichenberger <[email protected]>

tungld

Could you add into inference_backend.py backend tests for GroupNormalization?

tungld · 2023-10-31T04:37:13Z

src/Transform/ONNX/Decompose.cpp

+      assert(C % numGroups == 0 && "expected numGroups to divide C");
+      layerNormShapeVal.emplace_back(C / numGroups);
+    } else
+      layerNormShapeVal.emplace_back(-1);


should be ShapedType::kDynamic instead of -1? MLIR is no longer using -1 for dynamic dimension

tungld · 2023-10-31T04:41:53Z

src/Transform/ONNX/Decompose.cpp

+    Type biasScaleType = RankedTensorType::get(biasScaleShape, elementType);
+    Value newScale = create.onnx.unsqueeze(biasScaleType, scale, axes);
+    Value newBias = create.onnx.unsqueeze(biasScaleType, bias, axes);
+    // Convert input from N x C x D1...Dn to N x (NG x C/NG) x D1...Dn.


Do we need to handle special cases differently for the following cases? i.e. not need to split C into groups.

When the number of groups is the same as the number of channels, this operator is equivalent to InstanceNormalization. When there is only one group, this operator is equivalent to LayerNormalization.

or we finally get the same performance even though we split C?

From what I see, the consecutive reshape get optimized away, so there is no need to special case these two corner cases.

tungld · 2023-10-31T04:49:16Z

src/Dialect/ONNX/DialectBuilder.cpp

+  return IntegerAttr::get(b().getIntegerType(64, /*isSigned=*/true),
+      APInt(64, n, /*isSigned=*/true));
+}
+


Attribute is independent of ONNX. Perhaps, we should have a AttrBuilder that can be used with different dialects. This can be done by another PR.

This code was manually copied over and over again; got tired of it so I made a local private function. Would you like it to be in the MLIR dialect? If you want a new attribute builder, let's discuss which other functions you may want to have there.

yes, MLIR dialect looks ok.

@tungld, I decided against putting it into mlir builder, as other wise I need to build that other builder each time I simply want this. If you don't like this, I can put it as a static function (not part of the class) as a simple helper, or remove it all together.

OK, let's have another PR to deal with Attribute since this is not the main issue in this patch. Another candidate is something like Support/TypeUtilities.hpp that provides utility functions about type.

tungld · 2023-10-31T05:06:25Z

src/Conversion/ONNXToKrnl/NN/Normalization.cpp

+    PartialSpecified_FullySpecified, // Flattened to 2D.
+    FullySpecified_Scalar,           // Flattened to 2D.
+    FullySpecified_FullySpecified    // Flattened to 2D.
+  };


It's a worth extension, the output is more specific. Thanks!

sorenlassen

LGTM

I didn't review the ONNXToKrnl stuff but the rest looks good to me

sorenlassen · 2023-10-31T14:57:23Z

src/Transform/ONNX/Decompose.cpp

+      PatternRewriter &rewriter) const final {
+    // Match.
+    Value input = instanceNormOp.getInput();
+    if (!input.getType().isa<ShapedType>())


also fail if input has no rank, right?

good suggestion, included in the next batch of changes/

sorenlassen · 2023-10-31T15:11:49Z

src/Transform/ONNX/Decompose.cpp

+    auto inputShape = inputType.getShape();
+    int64_t C = inputShape[1];
+    int64_t inputRank = inputType.getRank();
+    assert(inputRank > 2 && "expected instance norm with input ranks > 2");


conceptually, is the "magic number" 2 here the same as the 2 in axis in line 775? consider naming the constant 2 in a way so that you can refer to it by name, both in InstanceNormIntoLayerNormPattern and GroupNormIntoLayerNormPattern where the number 2 appears multiple times too

sorenlassen · 2023-10-31T15:22:17Z

src/Dialect/ONNX/DialectBuilder.cpp

@@ -28,6 +28,11 @@ namespace onnx_mlir {

 //====-------------------------- ONNX Builder ---------------------------===//

+IntegerAttr OnnxBuilder::getSignedInt64Attr(int64_t n) const {
+  return IntegerAttr::get(b().getIntegerType(64, /*isSigned=*/true),
+      APInt(64, n, /*isSigned=*/true));


you can just pass n as 2nd arg to IntegerAttr::get, no need for APInt

nice, thanks.

AlexandreEichenberger · 2023-11-01T13:41:30Z

This PR uncovered issue #2601

…sage Signed-off-by: Alexandre Eichenberger <[email protected]>

Signed-off-by: Alexandre Eichenberger <[email protected]>

tungld

LGTM.

tungld · 2023-11-02T08:55:55Z

@AlexandreEichenberger does this issue #2601 block this patch? I am working on a fix for that issue, and it'll take a bit of time.

AlexandreEichenberger · 2023-11-02T12:32:11Z

does this issue #2601 block this patch?

It does, but there is no rush to this patch.

Signed-off-by: Alexandre Eichenberger <[email protected]>

jenkins-droid · 2023-11-09T22:30:20Z

Jenkins Linux s390x Build #13346 [push] Layernorm: convert insta... started at 17:30

jenkins-droid · 2023-11-09T22:30:22Z

Jenkins Linux ppc64le Build #12339 [push] Layernorm: convert insta... started at 17:38

jenkins-droid · 2023-11-09T22:30:22Z

Jenkins Linux amd64 Build #13321 [push] Layernorm: convert insta... started at 16:30

jenkins-droid · 2023-11-10T00:20:21Z

Jenkins Linux s390x Build #13346 [push] Layernorm: convert insta... passed after 1 hr 50 min

jenkins-droid · 2023-11-10T00:35:07Z

Jenkins Linux ppc64le Build #12339 [push] Layernorm: convert insta... passed after 2 hr 4 min

jenkins-droid · 2023-11-10T02:25:38Z

Jenkins Linux amd64 Build #13321 [push] Layernorm: convert insta... failed after 3 hr 55 min

* detect LayerNorm in presence of reciprocal and div of 1 (onnx#2609) Signed-off-by: Alexandre Eichenberger <[email protected]> * [NNPA] Use F16 as element type for zTensor (onnx#2611) * Use f16 as element type for zTensor Signed-off-by: Tung D. Le <[email protected]> --------- Signed-off-by: Tung D. Le <[email protected]> * Layernorm: convert instance norm and group norm to layer norm. (onnx#2595) Signed-off-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Tung D. Le <[email protected]> * Parse and set --mcpu in onnx-mlir-opt command (onnx#2614) Signed-off-by: Tung D. Le <[email protected]> * Update sqrt.mlir * Update sqrt.mlir * Update invsqrt.mlir * Update invsqrt.mlir * Update invsqrt.mlir * Update invsqrt.mlir Co-authored-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Tung D. Le <[email protected]> Co-authored-by: C-P2PN897 <[email protected]>

* detect LayerNorm in presence of reciprocal and div of 1 (onnx#2609) Signed-off-by: Alexandre Eichenberger <[email protected]> * [NNPA] Use F16 as element type for zTensor (onnx#2611) * Use f16 as element type for zTensor Signed-off-by: Tung D. Le <[email protected]> --------- Signed-off-by: Tung D. Le <[email protected]> * Layernorm: convert instance norm and group norm to layer norm. (onnx#2595) Signed-off-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Tung D. Le <[email protected]> * Parse and set --mcpu in onnx-mlir-opt command (onnx#2614) Signed-off-by: Tung D. Le <[email protected]> * Import dim_param for model inputs and outputs (onnx#2616) * Import dim_param for model inputs and outputs * use argument attributes Signed-off-by: Tung D. Le <[email protected]> --------- Signed-off-by: Tung D. Le <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> * [DialectBuilder] add builder funcrions for ONNXSumOp and ONNXConvOp (onnx#2572) The DialectBuilder class seems to be missing the function create the ONNXSumOp and ONNXConOp nodes and check their shape. This patch adds the necessary functions. Signed-off-by: Ashay Rane <[email protected]> Signed-off-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> * [StableHLO] Lowers PadOp (constant mode) & GatherElements Op to StableHLO (onnx#2602) * [Stablehlo] Pad constant mode & GatherElements to Stablehlo Signed-off-by: chongsong.chen <[email protected]> Signed-off-by: Yan Xu <[email protected]> Co-authored-by: chongsong.chen <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> * [build] Add cmake option to enable/disable Java components build (onnx#2613) * Add ONNX_MLIR_ENABLE_JAVA cmake option (default TRUE) Signed-off-by: Boyana Norris <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Alexandre Eichenberger <[email protected]> Co-authored-by: Tung D. Le <[email protected]> Co-authored-by: Ashay Rane <[email protected]> Co-authored-by: Yan Xu <[email protected]> Co-authored-by: chongsong.chen <[email protected]> Co-authored-by: Boyana Norris <[email protected]>

* 'main' of github.ibm.com:zosdev/onnx-mlir: Use dim_params in dynamic dimension analysis (onnx#2620) Update rapidcheck to include the fix for missing <cstdint> include (onnx#2623) Initial changes for llvm uplift (onnx#2568) [build] Add cmake option to enable/disable Java components build (onnx#2613) [StableHLO] Lowers PadOp (constant mode) & GatherElements Op to StableHLO (onnx#2602) [DialectBuilder] add builder funcrions for ONNXSumOp and ONNXConvOp (onnx#2572) Import dim_param for model inputs and outputs (onnx#2616) Parse and set --mcpu in onnx-mlir-opt command (onnx#2614) Layernorm: convert instance norm and group norm to layer norm. (onnx#2595) [NNPA] Use F16 as element type for zTensor (onnx#2611) detect LayerNorm in presence of reciprocal and div of 1 (onnx#2609) # Conflicts: # test/mlir/conversion/onnx_to_krnl/NN/Normalization_O3_SIMD_canonicalize.mlir

AlexandreEichenberger and others added 30 commits October 13, 2023 15:31

initial verification

706b5ca

Signed-off-by: Alexandre Eichenberger <[email protected]>

update

b29e4d9

shape inference for layer norm

f0acd49

Signed-off-by: Alexandre Eichenberger <[email protected]>

format

6f9a4a6

Signed-off-by: Alexandre Eichenberger <[email protected]>

Merge branch 'main' into layernorm

4ca8989

version that appears to work

2e233fe

Signed-off-by: Alexandre Eichenberger <[email protected]>

added backend tests

8c04f34

Signed-off-by: Alexandre Eichenberger <[email protected]>

updates

2e87037

Signed-off-by: Alexandre Eichenberger <[email protected]>

format

793abb9

Signed-off-by: Alexandre Eichenberger <[email protected]>

simplify onnx builder

af1f344

Signed-off-by: Alexandre Eichenberger <[email protected]>

format

4b6aaff

Signed-off-by: Alexandre Eichenberger <[email protected]>

Merge branch 'main' into layernorm

45df2e0

Merge branch 'main' into layernorm

8c05574

Merge branch 'main' into layernorm

3f91751

update

3a3dfaa

update

c8765c0

Signed-off-by: Alexandre Eichenberger <[email protected]>

update

8422de8

response to comments

c105881

Signed-off-by: Alexandre Eichenberger <[email protected]>

Merge branch 'layernorm' of https://github.com/AlexandreEichenberger/…

fb882ce

…onnx-mlir into layernorm

format

d5fa552

Signed-off-by: Alexandre Eichenberger <[email protected]>

Merge branch 'layernorm' into layernorm_v2

cf1d006

Merge branch 'main' into layernorm_v2

6404946

comments

9e88d90

Signed-off-by: Alexandre Eichenberger <[email protected]>

update

48fbbb0

Merge branch 'layernorm_v2' of https://github.com/AlexandreEichenberg…

e1d1ff5

…er/onnx-mlir into layernorm_v2

update

e18a309

support for flatten in MLIR builder

3ee663a

Signed-off-by: Alexandre Eichenberger <[email protected]>

progress

95d08a0

Signed-off-by: Alexandre Eichenberger <[email protected]>

update

b2198b5

migrated function that gen onnx code as an independent function

ccf6f53

Signed-off-by: Alexandre Eichenberger <[email protected]>

working for group norm too

af31cd6

Signed-off-by: Alexandre Eichenberger <[email protected]>

AlexandreEichenberger added 3 commits October 30, 2023 11:29

update

c930df2

format

f76ff92

Signed-off-by: Alexandre Eichenberger <[email protected]>

format

93860ff

Signed-off-by: Alexandre Eichenberger <[email protected]>

AlexandreEichenberger requested review from tungld and sorenlassen October 30, 2023 15:52

update

efa1f9d

tungld reviewed Oct 31, 2023

View reviewed changes

sorenlassen reviewed Oct 31, 2023

View reviewed changes

AlexandreEichenberger added 3 commits November 1, 2023 10:43

requested changes, plus adding op to refine to get a better error mes…

683adac

…sage Signed-off-by: Alexandre Eichenberger <[email protected]>

update

ea5f5ac

responce to comments

2056d3b

Signed-off-by: Alexandre Eichenberger <[email protected]>

tungld approved these changes Nov 2, 2023

View reviewed changes

AlexandreEichenberger added 3 commits November 6, 2023 15:11

update

fd3ac0b

Signed-off-by: Alexandre Eichenberger <[email protected]>

Merge branch 'main' into layernorm_v4

1dcb6a5

Merge branch 'main' into layernorm_v4

1a6c7be

AlexandreEichenberger merged commit 1c13ecf into onnx:main Nov 9, 2023
5 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Layernorm: convert instance norm and group norm to layer norm. #2595

Layernorm: convert instance norm and group norm to layer norm. #2595

AlexandreEichenberger commented Oct 30, 2023

AlexandreEichenberger commented Oct 30, 2023

tungld left a comment

tungld Oct 31, 2023

AlexandreEichenberger Oct 31, 2023

tungld Oct 31, 2023

AlexandreEichenberger Nov 1, 2023

tungld Oct 31, 2023

AlexandreEichenberger Oct 31, 2023

tungld Oct 31, 2023

AlexandreEichenberger Nov 1, 2023

tungld Nov 2, 2023

tungld Oct 31, 2023

sorenlassen left a comment

sorenlassen Oct 31, 2023

AlexandreEichenberger Nov 1, 2023

sorenlassen Oct 31, 2023

AlexandreEichenberger Nov 1, 2023

sorenlassen Oct 31, 2023

AlexandreEichenberger Nov 1, 2023

AlexandreEichenberger commented Nov 1, 2023

tungld left a comment

tungld commented Nov 2, 2023

AlexandreEichenberger commented Nov 2, 2023 •

edited

Loading

jenkins-droid commented Nov 9, 2023

jenkins-droid commented Nov 9, 2023

jenkins-droid commented Nov 9, 2023

jenkins-droid commented Nov 10, 2023

jenkins-droid commented Nov 10, 2023

jenkins-droid commented Nov 10, 2023

Layernorm: convert instance norm and group norm to layer norm. #2595

Layernorm: convert instance norm and group norm to layer norm. #2595

Conversation

AlexandreEichenberger commented Oct 30, 2023

AlexandreEichenberger commented Oct 30, 2023

tungld left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sorenlassen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

AlexandreEichenberger commented Nov 1, 2023

tungld left a comment

Choose a reason for hiding this comment

tungld commented Nov 2, 2023

AlexandreEichenberger commented Nov 2, 2023 • edited Loading

jenkins-droid commented Nov 9, 2023

jenkins-droid commented Nov 9, 2023

jenkins-droid commented Nov 9, 2023

jenkins-droid commented Nov 10, 2023

jenkins-droid commented Nov 10, 2023

jenkins-droid commented Nov 10, 2023

AlexandreEichenberger commented Nov 2, 2023 •

edited

Loading