[Mlir-commits] [mlir] [mlir][tosa] Remove `log(exp(x))` and `exp(log(x))` folders (PR #156434)

[Luke Hutton]

https://github.com/lhutton1 created https://github.com/llvm/llvm-project/pull/156434

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.

>From 1d52e781a664f9646491e86f4126100b09c0c9fb Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Tue, 2 Sep 2025 09:31:09 +0000
Subject: [PATCH] [mlir][tosa] Remove `log(exp(x))` and `exp(log(x))` folders

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.

Change-Id: Ieb9e897f8b2295f50fd9f50b9227707333c1fff3
---
 mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td  |  4 ----
 .../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 20 -------------------
 mlir/test/Dialect/Tosa/canonicalize.mlir      | 20 -------------------
 3 files changed, 44 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 416df6e87b11f..195f4cc57a197 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -1269,8 +1269,6 @@ def Tosa_ExpOp : Tosa_ElementwiseUnaryOp<"exp"> {
     Profile<[Tosa_PRO_FP]>,
     Extension<[Tosa_EXT_BF16]>,
   ];
-
-  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
@@ -1319,8 +1317,6 @@ def Tosa_LogOp : Tosa_ElementwiseUnaryOp<"log"> {
     Profile<[Tosa_PRO_FP]>,
     Extension<[Tosa_EXT_BF16]>,
   ];
-
-  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 42bce6931a002..35b8f53894797 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1555,26 +1555,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 b6d68c94f436b..4496b1eda1adf 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>