[Mlir-commits] [mlir] a7380fb - [mlir][tosa] Fold consecutive tosa.abs

[Kai Sasaki]

Author: Kai Sasaki
Date: 2023-05-22T15:00:05+09:00
New Revision: a7380fb702a69fd9fde3e6ce6107ed11449aa4ff

URL: https://github.com/llvm/llvm-project/commit/a7380fb702a69fd9fde3e6ce6107ed11449aa4ff
DIFF: https://github.com/llvm/llvm-project/commit/a7380fb702a69fd9fde3e6ce6107ed11449aa4ff.diff

LOG: [mlir][tosa] Fold consecutive tosa.abs

Consecutive tosa.abs can be fold as single abs operation since the second one has no impact.

Differential Revision: https://reviews.llvm.org/D150836

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
    mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
    mlir/test/Dialect/Tosa/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 842ac746eb664..f064e7a180441 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -944,6 +944,8 @@ def Tosa_AbsOp : Tosa_Op<"abs", [
   let results = (outs
     Tosa_Tensor:$output
   );
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 2477f154bc501..289f19559ad64 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1027,3 +1027,13 @@ OpFoldResult tosa::NegateOp::fold(FoldAdaptor adaptor) {
 
   return {};
 }
+
+OpFoldResult tosa::AbsOp::fold(FoldAdaptor adaptor) {
+  auto input = getInput1();
+  // Element-wise abs(abs(x)) = abs(x)
+  if (auto op = input.getDefiningOp<tosa::AbsOp>()) {
+    return input;
+  }
+
+  return {};
+}

diff  --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 3379850faf272..f7c3dbeb40d07 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -549,3 +549,13 @@ func.func @fold_negate_negate(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
   return %1 : tensor<?x1xf32>
 }
 
+// -----
+
+// CHECK-LABEL: @fold_abs_abs
+func.func @fold_abs_abs(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
+  // CHECK: %[[ABS:.*]] = "tosa.abs"(%arg{{.*}}) : (tensor<?x1xf32>) -> tensor<?x1xf32>
+  // CHECK: return %[[ABS]] : tensor<?x1xf32>
+  %0 = "tosa.abs"(%arg0) : (tensor<?x1xf32>) -> tensor<?x1xf32>
+  %1 = "tosa.abs"(%0) : (tensor<?x1xf32>) -> tensor<?x1xf32>
+  return %1 : tensor<?x1xf32>
+}