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

[Luke Hutton]

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

>From 84dae01be02aa21a9489df1b12d0f8d9330ab3a1 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 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>