TOSA legalization updates for spec v0.22, Part 1 #48193

sjarus · 2021-03-30T17:03:29Z

Updated gather, gather_nd and resize.
Fix precision of TFL avgpool2d, quantize. Add squareddifference
Add left/right shift, leaky_relu, one_hot
Update relu6/relu_n1_to1
Numerical precision in tf.fakequant legalization

Signed-off-by: Suraj Sudhir suraj.sudhir@arm.com
Change-Id: Idf23c3b24342f75ee7d1eb22dab6ffe27e8710b1

- Updated gather, gather_nd and resize. - Fix precision of TFL avgpool2d, quantize. Add squareddifference - Add left/right shift, leaky_relu, one_hot - Update relu6/relu_n1_to1 - Numerical precision in tf.fakequant legalization Signed-off-by: Suraj Sudhir <suraj.sudhir@arm.com> Change-Id: Idf23c3b24342f75ee7d1eb22dab6ffe27e8710b1

sjarus · 2021-03-30T17:08:51Z

@stellaraccident and @rsuderman , here are the part 1 set of updates to legalizations from TF/TFLite to TOSA aligned to the recent LLVM side update. That change was picked up by TensorFlow yesterday so we're pushing this out.

rsuderman · 2021-03-30T21:39:41Z

@stellaraccident and @rsuderman , here are the part 1 set of updates to legalizations from TF/TFLite to TOSA aligned to the recent LLVM side update. That change was picked up by TensorFlow yesterday so we're pushing this out.

Should be fine. We can revalidate internally when we land the CL.

rsuderman · 2021-03-30T21:43:57Z

tensorflow/compiler/mlir/tosa/BUILD

@@ -79,6 +79,7 @@ cc_library(
        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
        "//tensorflow/core:framework",
        "//tensorflow/core/kernels:conv_grad_shape_utils",
+        "//tensorflow/core/kernels:fake_quant_ops",


I'm a little confused why this is added. It feels unrelated to the rest of the CL.

It implements tensorflow::Nudge that is invoked from convertFakeQuantOp() in legalize_common.cc

Thanks - i will try to review in detail tonight but have a lot in my queue and it may slip to tomorrow.

This will add a cross dependency on a large body of TensorFlow to the TFLite path (est. 1000 additional source files?). I would rather just copy the Nudge function locally. (this has been done a couple of times for dependency reasons already)

Thanks for suggestion. Will try to replace it with a local implementation in next round.

stellaraccident

Thank you for the contributions. There are a number of comments ranging from local/style, to design to testing. I think the main design point has to do with the override_zero_point attribute and doing that in a way that does not introduce a cross pass dependency. The other big request is upgraded test coverage. I'm happy to jump on some kind of shared space and help educate on some practices here. In addition, I'd suggest looking at the way that some of the legalize_tf tests are done in the tensorflow/xla side for an idea of more of what we are looking for.

stellaraccident · 2021-04-01T03:51:42Z

tensorflow/compiler/mlir/tosa/transforms/convert_tfl_uint8.cc

+    auto tfl_avgpool2d_op = cast<TFL::AveragePool2DOp>(op);
+
+    auto input_type =
+        tfl_avgpool2d_op.input().getType().dyn_cast<RankedTensorType>();


input_type is constrained by the op definition to be RankedTensorType, right? If that is the case, why dyn_cast (you aren't actually checking if it is !nullptr). Suggest either using cast if this is an invariant or continuing to use dyn_cast and early returning if any dyn_cast returns nullptr.

My bad. Will fix in next round.

stellaraccident · 2021-04-01T03:52:29Z

tensorflow/compiler/mlir/tosa/transforms/convert_tfl_uint8.cc

@@ -128,6 +128,48 @@ struct ConvertUint8QConstOp : public RewritePattern {
  }
 };

+struct ConvertAveragePool2DOp : public RewritePattern {


Why isn't this an OpRewritePatternTFL::AveragePool2DOp? This would then simplify your constructor and let you remove the cast in the first line of matchAndRewrite.

Do you have example for this that we can dumbly follow?

stellaraccident · 2021-04-01T03:53:34Z

tensorflow/compiler/mlir/tosa/transforms/convert_tfl_uint8.cc

+
+    // Annotate attribute to match TFLite average pool2d rounding behavior.
+    auto key =
+        mlir::Identifier::get(kOverrideZeropointAttrName, builder.getContext());


Prefer builder.getIdentifier(...)

Thanks. Will update in next round.

stellaraccident · 2021-04-01T03:58:10Z

tensorflow/compiler/mlir/tosa/transforms/convert_tfl_uint8.cc

+    auto value = builder.getI32IntegerAttr(override_zeropoint);
+    op->setAttr(key, value);
+
+    return success();


I'm still reading how this fits together, but as written, this pattern sketches me out because it is mutating in place and doesn't seem to either produce a new op or alter its match behavior based on what gets set. I don't think this is going to have the desired effect. Will try to discern the goal and offer a suggestion elsewhere.

I can't quite see here why you need an in-place pattern here (versus a helper at the point of transformation which checks these conditions and returns the zero point override).

Please see my comment on top.

stellaraccident · 2021-04-01T04:04:51Z

tensorflow/compiler/mlir/tosa/transforms/legalize_tfl.cc

+  if (input_type.getElementType().isa<UniformQuantizedType>()) {
+    // Search for attribute annotated by --tosa-convert-tfl-uint8 pass.
+    // This is needed in average pool2d to match TFLite rounding behavior.
+    IntegerAttr override_zeropoint_attr =


These kind of cross-pass action at a distance attribute dependencies need to be avoided. Why can't you have a helper which performs the check in the pattern which sets this attribute and just call it here (i.e. a static calculateZeropointOverride(TFL::AveragePool2DOp op)).

Even better you could define the above to return an IntegerAttr conditionally if the input_type is quantized (or nullptr otherwise), then collapse all of the branching in this transform.

Explained in the comment on top. QU8 and QI8 is not distinguishable after convert_tfl_uint8 pass.

This piece appears to be the largest blocker for this update. We're aware - as Kevin mentioned - that the approach isn't ideal. Is it ok if we get all other changes done, then do a separate update with a cleaner implementation of this cross pass issue ?

stellaraccident · 2021-04-01T04:44:10Z

tensorflow/compiler/mlir/tosa/transforms/legalize_common.cc

+  for (int32_t i = 0; i < indices_type.getRank(); i++) {
+    int32_t dim = indices_type.getShape()[i];
+    N *= dim;
+    if (i >= axis)


Unlike LLVM style, google style wants braces around these.

Will add braces in next round.

stellaraccident · 2021-04-01T04:44:39Z

tensorflow/compiler/mlir/tosa/transforms/legalize_common.cc

+    axis = indices_type.getRank();
+  }
+
+  int32_t N = 1, W = 1, C = 1;


Prefer just int instead of int32_t unless if there is a strong reason (here and below).

Thanks will update in next round.

stellaraccident · 2021-04-01T04:48:08Z

tensorflow/compiler/mlir/tosa/tests/tf-to-tosa-pipeline.mlir

+// CHECK: tosa.reshape
+// CHECK: tosa.transpose
+// CHECK: tosa.reshape
+func @test_one_hot(%arg0: tensor<4x4xi32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<4x4x2xf32> {


General note on tests (and specific note on this): These are non trivial lowerings and we should have more explicit tests. The tests do not need to verify every detail of the patterns but should minimally capture SSA values and constants. Getting the testing granularity right is a bit of an art, so if you want help on this, we can bust some examples out to a doc or gist or something and work on it together if that is helpful (been meaning to ask for test upgrades for some time and realize there may be some education here).

Ah so we can check the SSA values as well. Didn't realize we can do that.
I'd assume all we need to do is append that to "// CHECK: tosa.XXX", for example like "// CHECK: tosa.XXX(%0)"?
or it's not that simple as I expect?

You can see similar examples in the tensorflow to hlo lowering tests:

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/xla/tests/legalize-tf.mlir#L64

The problem with not validating the large sequences is incorrect constants or changing which values are passed to which ops.

Nit: You can drop all the shape information in the line to make things more succinct.

I.e. The following line:

// CHECK: %[[VAR0:.*]] = "tosa.const"() {value = dense<[0, 2, 1]> : tensor<3xi32>} : () -> tensor<3xi32>

Can be simplified to

// CHECK: %[[VAR0:.*]] = "tosa.const"() {value = dense<[0, 2, 1]> : tensor<3xi32>}

Generally we want to make sure check tests are mostly a sequence of correct values.

stellaraccident · 2021-04-01T04:49:06Z

tensorflow/compiler/mlir/tosa/tests/tf-to-tosa-pipeline.mlir

@@ -105,9 +105,16 @@ func @test_relu6(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
 // -----



I didn't double check, but I think that some of the op conversions added are missing tests (gather?). Please audit test coverage for both ops and critical branches in the patterns.

Thanks. Will fix in next round.

Still missing tests for gather support. Because we are adding functionality for a gather lowering we will need a test to validate the lowering is working correctly.

stellaraccident · 2021-04-01T04:53:11Z

tensorflow/compiler/mlir/tosa/BUILD

@@ -79,6 +79,7 @@ cc_library(
        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
        "//tensorflow/core:framework",
        "//tensorflow/core/kernels:conv_grad_shape_utils",
+        "//tensorflow/core/kernels:fake_quant_ops",


This will add a cross dependency on a large body of TensorFlow to the TFLite path (est. 1000 additional source files?). I would rather just copy the Nudge function locally. (this has been done a couple of times for dependency reasons already)

armkevincheng · 2021-04-06T17:33:30Z

I agree the cross-pass design is bad, but that was the only way I thought I can achieve what I want.

Briefly explain what I'm doing here:
We're trying to match the rounding behavior to TFL::AveragePool2DOp, which has different rounding behavior between QI8 and QU8 (since it's rounding on its storage type/range).

In the TFL to TOSA pass pipeline, the convert_tfl_uint8 pass is called first, and converts all the QU8 into QI8, so the pass after, e.g. legalize_tfl, only need to deal with QI8, but at the same time can't distinguish between those two cases.

The trick we played here (which is bad I agree) is to annotate the override_zeropoint attribute based on if it's QU8 or QI8 before QU8 is converted. When pass reaches legalize_tfl, it checks if such attribute exists. If it does, then we override that zeropoint with what's stored in the attribute, and we build TOSA::AveragePool2dOp with it.

Could something you mentioned above solve the problem without using the cross-pass design?

armkevincheng · 2021-04-06T17:35:33Z

Again thank for all the feedbacks. They're all pretty helpful and we'll prepare next round of review as soon as possible.

- Use tensorflow::Nudge implementation locally to avoid build dependency - Cleanup in coding artifacts - Disable cross-pass AvgPool2d int8/uint8 handling until better solution Change-Id: I1e0421ee336e4e137aebfa2856870fa3838169e6 Signed-off-by: Suraj Sudhir <suraj.sudhir@arm.com>

sjarus · 2021-04-09T17:14:59Z

We've removed avgpool until the cross pass dependency is resolved, but have updated other things. We hope we can upstream this, and follow up the remaining pieces in a separate PR related to the remaining TOSA v0.22 changes still to follow.

rsuderman · 2021-04-12T18:03:33Z

tensorflow/compiler/mlir/tosa/tests/tf-to-tosa-pipeline.mlir

+// CHECK: tosa.reshape
+// CHECK: tosa.transpose
+// CHECK: tosa.reshape
+func @test_one_hot(%arg0: tensor<4x4xi32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<4x4x2xf32> {


You can see similar examples in the tensorflow to hlo lowering tests:

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/xla/tests/legalize-tf.mlir#L64

The problem with not validating the large sequences is incorrect constants or changing which values are passed to which ops.

rsuderman · 2021-04-12T18:04:32Z

tensorflow/compiler/mlir/tosa/tests/tf-to-tosa-pipeline.mlir

@@ -105,9 +105,16 @@ func @test_relu6(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
 // -----



Still missing tests for gather support. Because we are adding functionality for a gather lowering we will need a test to validate the lowering is working correctly.

rsuderman · 2021-04-12T18:06:42Z

tensorflow/compiler/mlir/tosa/transforms/legalize_tfl.cc

  if (!element_type) return failure();

+  // In some cases output_type is dynamic shape as tensor<*xelement_type>


This kinda of case usually occurs when tensorflow's shape propagation is not working correctly. Passes should assume valid shape propagation as it is an assumed-correct upstream behavior.

rsuderman · 2021-04-12T18:07:41Z

tensorflow/compiler/mlir/tosa/transforms/legalize_tf.cc

+    alpha = tmpAttr.getValueAsDouble();
+  }
+
+  Value const_zero = getTosaConstTensorSingleF32(rewriter, op, 0.0);


Sounds fine.

sjarus · 2021-04-12T21:38:05Z

We're updating test_one_hot to follow the suggested approach. If that looks good, we'll follow up with a clean up of the entire set of tests in a follow up PR since it is an independent task. However we intend to add the gather tests to this PR, hopefully enabling us to close this one out.

Fixed tfl.quantize legalization output_type setup. Change-Id: I4a4826f653941299916229d8e2d0e094342f8ef4 Signed-off-by: Suraj Sudhir <suraj.sudhir@arm.com>

rsuderman

Last two comments related to tests. Once these wrap we should be good to land.

rsuderman · 2021-04-13T19:52:37Z

tensorflow/compiler/mlir/tosa/tests/tf-to-tosa-pipeline.mlir

+// CHECK: tosa.reshape
+// CHECK: tosa.transpose
+// CHECK: tosa.reshape
+func @test_one_hot(%arg0: tensor<4x4xi32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<4x4x2xf32> {


Nit: You can drop all the shape information in the line to make things more succinct.

I.e. The following line:

// CHECK: %[[VAR0:.*]] = "tosa.const"() {value = dense<[0, 2, 1]> : tensor<3xi32>} : () -> tensor<3xi32>

Can be simplified to

// CHECK: %[[VAR0:.*]] = "tosa.const"() {value = dense<[0, 2, 1]> : tensor<3xi32>}

Generally we want to make sure check tests are mostly a sequence of correct values.

rsuderman · 2021-04-13T19:54:46Z

tensorflow/compiler/mlir/tosa/tests/tf-to-tosa-pipeline.mlir

+// CHECK: tosa.reshape
+// CHECK: tosa.mul
+// CHECK: tosa.reshape
+// CHECK: tosa.add
 func @test_fakequant_with_min_max_args(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {


Perform the same value validation on the rest of the added tests. If they are short (e.g. less than 3-4 ops) its fine without them but most of these tests are pretty complex and we want to avoid future errors.

The plan is to clean up all the tests in a separate PR. Is that a workable option ? We tried the test_one_hot to get feedback on the right way to go about it, rather than change everything, then potentially have to do a second pass at everything.

sjarus · 2021-04-13T20:56:22Z

Please let us know if we can update all the TF + TFL tests in a separate PR, using the suggested template of test_one_hot . Since this is largely independent of this PR and the ones to follow it, we can parallelize the test update work.

rsuderman

We can land under the guarantee that the followup appears reasonably soon. We tend to avoid landing tests with brittle checks in case TF canonicalizations are updated. This can cause some unexpected failures.

sjarus · 2021-04-14T21:22:38Z

We can land under the guarantee that the followup appears reasonably soon. We tend to avoid landing tests with brittle checks in case TF canonicalizations are updated. This can cause some unexpected failures.

Current estimate is for the updated tests for both TF and TFL legalization to constitute a new PR within the next week. Is that reasonable ?

tensorflow/compiler/mlir/tosa/transforms/legalize_tfl.cc

tensorflow/compiler/mlir/tosa/transforms/legalize_common.cc

rsuderman · 2021-04-14T21:30:21Z

Found a set of changes required to fix some internal failures, mostly unused variable issues.

rsuderman · 2021-04-14T21:47:40Z

Working on landing this internally patching the changes in myself.

sjarus · 2021-04-14T21:48:33Z

Working on landing this internally patching the changes in myself.

Oops didn't see this. Thanks for handling this!

google-ml-butler bot added the size:XL CL Change Size:Extra Large label Mar 30, 2021

google-cla bot added the cla: yes label Mar 30, 2021

google-ml-butler bot requested a review from joker-eph March 30, 2021 17:03

rsuderman reviewed Mar 30, 2021

View reviewed changes

gbaned self-assigned this Mar 31, 2021

gbaned added this to Assigned Reviewer in PR Queue via automation Mar 31, 2021

stellaraccident reviewed Apr 1, 2021

View reviewed changes

gbaned added the awaiting review Pull request awaiting review label Apr 8, 2021

sjarus requested review from stellaraccident and rsuderman April 9, 2021 23:22

rsuderman suggested changes Apr 12, 2021

View reviewed changes

PR Queue automation moved this from Assigned Reviewer to Reviewer Requested Changes Apr 12, 2021

Added tests for legalizations in tensorflow#48193

0c29763

Fixed tfl.quantize legalization output_type setup. Change-Id: I4a4826f653941299916229d8e2d0e094342f8ef4 Signed-off-by: Suraj Sudhir <suraj.sudhir@arm.com>

sjarus requested a review from rsuderman April 12, 2021 22:10

rsuderman suggested changes Apr 13, 2021

View reviewed changes

sjarus requested a review from rsuderman April 14, 2021 16:21

rsuderman approved these changes Apr 14, 2021

View reviewed changes

google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Apr 14, 2021

PR Queue automation moved this from Reviewer Requested Changes to Approved by Reviewer Apr 14, 2021

kokoro-team removed the kokoro:force-run Tests on submitted change label Apr 14, 2021

sjarus requested a review from rsuderman April 14, 2021 21:22

rsuderman suggested changes Apr 14, 2021

View reviewed changes

PR Queue automation moved this from Approved by Reviewer to Reviewer Requested Changes Apr 14, 2021

copybara-service bot merged commit d9bcf21 into tensorflow:master Apr 14, 2021

PR Queue automation moved this from Reviewer Requested Changes to Merged Apr 14, 2021

armkevincheng mentioned this pull request Apr 16, 2021

Update TF/TFL -> TOSA legalization tests with SSA values associated #48573

Merged

		@@ -105,9 +105,16 @@ func @test_relu6(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
		// -----

		if (!element_type) return failure();

		// In some cases output_type is dynamic shape as tensor<*xelement_type>

TOSA legalization updates for spec v0.22, Part 1 #48193

TOSA legalization updates for spec v0.22, Part 1 #48193

Conversation

sjarus commented Mar 30, 2021

sjarus commented Mar 30, 2021

rsuderman commented Mar 30, 2021

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

stellaraccident 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

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

armkevincheng commented Apr 6, 2021

armkevincheng commented Apr 6, 2021

sjarus commented Apr 9, 2021

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

sjarus commented Apr 12, 2021

rsuderman 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

sjarus commented Apr 13, 2021

rsuderman left a comment

Choose a reason for hiding this comment

sjarus commented Apr 14, 2021

rsuderman commented Apr 14, 2021

rsuderman commented Apr 14, 2021

sjarus commented Apr 14, 2021