From 62ccc50613a6912c4e450046c5fcd710007e3f8e Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe@gmail.com>
Date: Wed, 3 May 2023 14:05:59 +0900
Subject: [PATCH] [mlir][tosa] Fold log(exp) to no-op

Element-wise log(exp) does no operation so that we can fold it into no-op effectively.

Reviewed By: eric-k256

Differential Revision: https://reviews.llvm.org/D149632
---
 mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td       |  2 ++
 mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp | 10 ++++++++++
 mlir/test/Dialect/Tosa/canonicalize.mlir           | 10 ++++++++++
 3 files changed, 22 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index e36ab18..015be1f 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -1043,6 +1043,8 @@ def Tosa_LogOp : Tosa_Op<"log", [
   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 c6851d9..7e161b12 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -996,3 +996,13 @@ OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
 
   return getInput1();
 }
+
+OpFoldResult tosa::LogOp::fold(FoldAdaptor adaptor) {
+  auto input = getInput1();
+  // Element-wise log(exp(x)) = x
+  if (auto op = input.getDefiningOp<tosa::ExpOp>()) {
+    return op.getInput1();
+  }
+
+  return {};
+}
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index adcb16d..9633a94 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -505,3 +505,13 @@ func.func @canonicalize_concat_slice_on_non_concat_axis(%arg0 : tensor<1x12x12xf
   %2 = "tosa.slice"(%0) {size = array<i64: 1, 3, 12>, start = array<i64: 1, 3, 12>} : (tensor<1x12x24xf32>) -> tensor<1x3x12xf32>
   return %1, %2 : tensor<1x6x12xf32>, tensor<1x3x12xf32>
 }
+
+// -----
+
+// CHECK-LABEL
+func.func @fold_log_exp(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
+  // CHECK: return %arg{{.*}} : tensor<?x1xf32>
+  %0 = "tosa.exp"(%arg0) : (tensor<?x1xf32>) -> tensor<?x1xf32>
+  %1 = "tosa.log"(%0) : (tensor<?x1xf32>) -> tensor<?x1xf32>
+  return %1 : tensor<?x1xf32>
+}
-- 
2.7.4