[tosa] Legalize tfl.atan2 to tosa.table operation #59100

lhutton1 · 2023-01-04T22:12:37Z

In a similar fashion to tfl.sin and tlf.cos legalization, a tosa.table provides an atan approximation. Additional logic is then used to determine the correct quadrant of the atan2 function.

Signed-off-by: Luke Hutton luke.hutton@arm.com

jpienaar

Forgot to hit send on comments ...

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

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

lhutton1

Thanks for the review @jpienaar! I left some responses below

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

gbaned · 2023-01-11T05:27:15Z

Hi @lhutton1 Can you please resolve conflicts? Thank you!

lhutton1 · 2023-01-17T16:33:15Z

Friendly ping @jpienaar

jpienaar

Thanks for the ping

jpienaar · 2023-01-17T16:36:35Z

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

@@ -640,6 +640,47 @@ func.func @test_cos(%arg0: tensor<10xf32>) -> tensor<*xf32> {

 // -----

+// CHECK-LABEL: test_atan2
+// CHECK-SAME: -> tensor<13x21x3xf32>
+// CHECK: %[[VAL_2:.*]] = "tosa.const"() {value = dense<2.000000e+00> : tensor<1x1x1xf32>} : () -> tensor<1x1x1xf32>


Prefer CHECK-DAG for constants (there is an internal check for this that ends up blocking presubmits, but seems not upstream)

done thanks

jpienaar · 2023-01-17T16:37:46Z

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

+
+  Type input_y_ety = input_y_ty.getElementType();
+  Type input_x_ety = input_x_ty.getElementType();
+  Type output_ety = output_ty.getElementType();


Check for output_ty above too to avoid null here or just use cast above - I don't think TFL::Atan2Op is legal unless it produces a ShapedType

Oops seems I left these in after the last review - removed them

jpienaar · 2023-01-17T16:40:16Z

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

+  auto tfl_atan2_op = cast<TFL::Atan2Op>(op);
+  Location loc = op->getLoc();
+  Value input_y = tfl_atan2_op.getY();
+  RankedTensorType input_y_ty = input_y.getType().dyn_cast<RankedTensorType>();


With the new casting mechanisms in LLVM one can now finally do

auto input_y_ty = dyn_cast(input_y.getType());

prefer to not repeat type for these (for those that are outputs of cast's we use auto as type is clear locally).

Thanks, looks cleaner. I might have misunderstood the dyn_cast(...) suggestion here, but I was only able to get it working with a templated type dyn_cast<RankedTensorType>(...)

jpienaar · 2023-01-17T16:40:32Z

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

+  Type input_x_ety = input_x_ty.getElementType();
+  Type output_ety = output_ty.getElementType();
+
+  bool op_is_fp = input_y_ty.getElementType().isF32();


Use input_y_ety?

Removed as they are only used once anyway

jpienaar · 2023-01-17T16:41:36Z

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

+
+  if (!input_y_ty || !input_x_ty) {
+    return rewriter.notifyMatchFailure(
+        op, "ConvertTFLAtan2Op: ranked inputs required");


Update these to keep it consistent with the others (skipping the pattern name in error message)

Done - apologies, I realized I missed this from your previous review

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

jpienaar · 2023-01-17T16:45:12Z

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

+
+  // 2. Scale and translate the normalized domain to the table domain. This
+  // includes a translating and scaling to [-int16_max, int16_max] and casting
+  // to an i16.


Add reason for casting here too.

jpienaar · 2023-01-17T16:45:40Z

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

+      rewriter, loc, output_ty.clone(rewriter.getIntegerType(32)), casted,
+      table_const);
+
+  // 4. The range of table is a 23-bit two's compliment value. Normalize the


Fixed, thanks!

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

In a similar fashion to tfl.sin and tlf.cos legalization, a tosa.table provides an atan approximation. Additional logic is then used to determine the correct quadrant of the atan2 function. Signed-off-by: Luke Hutton <luke.hutton@arm.com> Change-Id: Iae1384009d825d01e5cf48ad7c3ff8fba77114cf

* remove unnecessary type checks * add note about numerical behaviour of std::atan2 * improve error message for expected inputs * undo change updating copyright year Change-Id: Iea8339da437a5ff3e6fe065c715c2c97e696fdbb

* Cleanup of error messages. * Cleanup casting inputs/outputs. * Use CHECK-DAG instead of CHECK for constants. * Spelling Change-Id: Ibf01ac5b711944bb2efbd7184820a304cbae8501 Signed-off-by: Luke Hutton <luke.hutton@arm.com>

lhutton1

Apologies for the late follow-up on this and thanks for the review @jpienaar, please find my replies below

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

lhutton1 · 2023-01-23T17:39:21Z

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

+      rewriter, loc, output_ty.clone(rewriter.getIntegerType(32)), casted,
+      table_const);
+
+  // 4. The range of table is a 23-bit two's compliment value. Normalize the


Fixed, thanks!

lhutton1 · 2023-01-23T18:30:09Z

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

+
+  Type input_y_ety = input_y_ty.getElementType();
+  Type input_x_ety = input_x_ty.getElementType();
+  Type output_ety = output_ty.getElementType();


Oops seems I left these in after the last review - removed them

lhutton1 · 2023-01-23T18:32:02Z

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

+  Type input_x_ety = input_x_ty.getElementType();
+  Type output_ety = output_ty.getElementType();
+
+  bool op_is_fp = input_y_ty.getElementType().isF32();


Removed as they are only used once anyway

lhutton1 · 2023-01-23T18:37:53Z

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

+  auto tfl_atan2_op = cast<TFL::Atan2Op>(op);
+  Location loc = op->getLoc();
+  Value input_y = tfl_atan2_op.getY();
+  RankedTensorType input_y_ty = input_y.getType().dyn_cast<RankedTensorType>();


Thanks, looks cleaner. I might have misunderstood the dyn_cast(...) suggestion here, but I was only able to get it working with a templated type dyn_cast<RankedTensorType>(...)

lhutton1 · 2023-01-23T18:38:05Z

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

+
+  if (!input_y_ty || !input_x_ty) {
+    return rewriter.notifyMatchFailure(
+        op, "ConvertTFLAtan2Op: ranked inputs required");


Done - apologies, I realized I missed this from your previous review

lhutton1 · 2023-01-23T18:40:39Z

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

+
+  // 2. Scale and translate the normalized domain to the table domain. This
+  // includes a translating and scaling to [-int16_max, int16_max] and casting
+  // to an i16.


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

lhutton1 · 2023-01-24T09:52:01Z

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

@@ -640,6 +640,47 @@ func.func @test_cos(%arg0: tensor<10xf32>) -> tensor<*xf32> {

 // -----

+// CHECK-LABEL: test_atan2
+// CHECK-SAME: -> tensor<13x21x3xf32>
+// CHECK: %[[VAL_2:.*]] = "tosa.const"() {value = dense<2.000000e+00> : tensor<1x1x1xf32>} : () -> tensor<1x1x1xf32>


done thanks

jpienaar

LG, check on consistency part and if the behavior TOSA side in case there were 0's but beyond that good. (and I saw your responses later, don't understand github PRs some days)

lhutton1 · 2023-01-26T23:03:41Z

Thanks! It was actually me who forgot to hit send 😅

google-ml-butler bot added the size:M CL Change Size: Medium label Jan 4, 2023

google-ml-butler bot assigned gbaned Jan 4, 2023

gbaned requested a review from jpienaar January 5, 2023 15:43

google-ml-butler bot added the awaiting review Pull request awaiting review label Jan 5, 2023

jpienaar reviewed Jan 6, 2023

View reviewed changes

lhutton1 commented Jan 6, 2023

View reviewed changes

tensorflow/compiler/mlir/tosa/transforms/legalize_tfl.cc Outdated Show resolved Hide resolved

tensorflow/compiler/mlir/tosa/transforms/legalize_tfl.cc Show resolved Hide resolved

lhutton1 force-pushed the tosa-atan2-tfl-legalization branch from b633a12 to 499bf34 Compare January 9, 2023 19:57

gbaned added stat:awaiting response Status - Awaiting response from author and removed awaiting review Pull request awaiting review labels Jan 11, 2023

lhutton1 force-pushed the tosa-atan2-tfl-legalization branch from 499bf34 to bdd550e Compare January 11, 2023 09:53

gbaned added awaiting review Pull request awaiting review and removed stat:awaiting response Status - Awaiting response from author labels Jan 12, 2023

jpienaar reviewed Jan 17, 2023

View reviewed changes

gbaned removed the awaiting review Pull request awaiting review label Jan 18, 2023

lhutton1 added 3 commits January 26, 2023 20:04

address comments

cbe054a

* remove unnecessary type checks * add note about numerical behaviour of std::atan2 * improve error message for expected inputs * undo change updating copyright year Change-Id: Iea8339da437a5ff3e6fe065c715c2c97e696fdbb

address comments

3986b07

* Cleanup of error messages. * Cleanup casting inputs/outputs. * Use CHECK-DAG instead of CHECK for constants. * Spelling Change-Id: Ibf01ac5b711944bb2efbd7184820a304cbae8501 Signed-off-by: Luke Hutton <luke.hutton@arm.com>

lhutton1 force-pushed the tosa-atan2-tfl-legalization branch from bdd550e to 3986b07 Compare January 26, 2023 22:40

lhutton1 commented Jan 26, 2023

View reviewed changes

jpienaar approved these changes Jan 26, 2023

View reviewed changes

google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Jan 26, 2023

kokoro-team removed the kokoro:force-run Tests on submitted change label Jan 26, 2023

copybara-service bot merged commit 0e89f1a into tensorflow:master Jan 28, 2023

lhutton1 deleted the tosa-atan2-tfl-legalization branch January 28, 2023 16:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[tosa] Legalize tfl.atan2 to tosa.table operation #59100

[tosa] Legalize tfl.atan2 to tosa.table operation #59100

lhutton1 commented Jan 4, 2023

jpienaar left a comment

lhutton1 left a comment

gbaned commented Jan 11, 2023

lhutton1 commented Jan 17, 2023

jpienaar left a comment

jpienaar Jan 17, 2023

lhutton1 Jan 24, 2023

jpienaar Jan 17, 2023

lhutton1 Jan 23, 2023

jpienaar Jan 17, 2023

lhutton1 Jan 23, 2023

jpienaar Jan 17, 2023

lhutton1 Jan 23, 2023

jpienaar Jan 17, 2023

lhutton1 Jan 23, 2023

jpienaar Jan 17, 2023

lhutton1 Jan 23, 2023

jpienaar Jan 17, 2023

lhutton1 Jan 23, 2023

lhutton1 left a comment

lhutton1 Jan 23, 2023

lhutton1 Jan 23, 2023

lhutton1 Jan 23, 2023

lhutton1 Jan 23, 2023

lhutton1 Jan 23, 2023

lhutton1 Jan 23, 2023

lhutton1 Jan 24, 2023

jpienaar left a comment

lhutton1 commented Jan 26, 2023

[tosa] Legalize tfl.atan2 to tosa.table operation #59100

[tosa] Legalize tfl.atan2 to tosa.table operation #59100

Conversation

lhutton1 commented Jan 4, 2023

jpienaar left a comment

Choose a reason for hiding this comment

lhutton1 left a comment

Choose a reason for hiding this comment

gbaned commented Jan 11, 2023

lhutton1 commented Jan 17, 2023

jpienaar 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

lhutton1 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

jpienaar left a comment

Choose a reason for hiding this comment

lhutton1 commented Jan 26, 2023