[mlir][tosa] Remove log(exp(x)) and exp(log(x)) folders

[lhutton1]

Both of the folders would reduce to x. These folders were vulnerable to overflow / underflow issues, resulting in a difference in numerical behaviour when running or not running the folders. For now they have been removed. We can consider restoring these as part of an optional and separate "fast-math" style of transformation in the future.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tosa

Author: Luke Hutton (lhutton1)

Changes

Both of the folders would reduce to x. These folders were vulnerable to overflow / underflow issues, resulting in a difference in numerical behaviour when running or not running the folders. For now they have been removed. We can consider restoring these as part of an optional and separate "fast-math" style of transformation in the future.

---

Full diff: https://github.com/llvm/llvm-project/pull/156434.diff

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td (-4)
• (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (-20)
• (modified) mlir/test/Dialect/Tosa/canonicalize.mlir (-20)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 953e7c304da85..48759f2a3c9e8 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -1331,8 +1331,6 @@ def Tosa_ExpOp : Tosa_ElementwiseUnaryOp<"exp"> {
     Extension<[Tosa_EXT_BF16]>,
   ];
 
-  let hasFolder = 1;
-
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
 }
 
@@ -1385,8 +1383,6 @@ def Tosa_LogOp : Tosa_ElementwiseUnaryOp<"log"> {
     Extension<[Tosa_EXT_BF16]>,
   ];
 
-  let hasFolder = 1;
-
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
 }
 
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 8d636460c667e..caf80165fc640 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1562,26 +1562,6 @@ 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 {};
-}
-
-OpFoldResult tosa::ExpOp::fold(FoldAdaptor adaptor) {
-  auto input = getInput1();
-  // Element-wise exp(log(x)) = x
-  if (auto op = input.getDefiningOp<tosa::LogOp>()) {
-    return op.getInput1();
-  }
-
-  return {};
-}
-
 OpFoldResult tosa::NegateOp::fold(FoldAdaptor adaptor) {
   // Element-wise negate(negate(x)) = x
   // iff all zero points are constant 0
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index fd2a3f1d361eb..e8525a5d2ed62 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1104,26 +1104,6 @@ func.func @canonicalize_pad_slice_dynamic_noupdate(%arg0: tensor<1x16x?x3xf32>)
 
 // -----
 
-// CHECK-LABEL: @fold_log_exp
-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>
-}
-
-// -----
-
-// CHECK-LABEL: @fold_exp_log
-func.func @fold_exp_log(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
-  // CHECK: return %arg{{.*}} : tensor<?x1xf32>
-  %0 = tosa.log %arg0 : (tensor<?x1xf32>) -> tensor<?x1xf32>
-  %1 = tosa.exp %0 : (tensor<?x1xf32>) -> tensor<?x1xf32>
-  return %1 : tensor<?x1xf32>
-}
-
-// -----
-
 // CHECK-LABEL: @fold_negate_negate
 func.func @fold_negate_negate(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
   // CHECK: return %arg{{.*}} : tensor<?x1xf32>

[lhutton1]

For context, please see: https://discourse.llvm.org/t/question-about-canonicalization-tensor-log-tensor-exp-x-x/88018

[Tai78641]

LGTM