[mlir][tosa] Fold log(exp) to no-op

[Lewuathe]

Summary

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

[Lewuathe]

Rebase

[eric-k256]

looks good to me.

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

def Tosa_LogOp : Tosa_Op<"log", [
 
   let arguments = (ins
     Tosa_Tensor:$input1
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: logical_not
 //===----------------------------------------------------------------------===//
 def Tosa_LogicalNotOp : Tosa_Op<"logical_not", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
   if (getConstantPerms(perms).failed())
     return {};
 
   if (!llvm::equal(llvm::seq<int64_t>(0, perms.size()), perms))
     return {};
 
   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 {};
+}

mlir/test/Dialect/Tosa/canonicalize.mlir

 // CHECK: %[[VAL_3:.*]] = "tosa.slice"(%[[VAL_1]]) <{size = array<i64: 1, 3, 12>, start = array<i64: 1, 3, 12>}> : (tensor<1x12x12xf32>) -> tensor<1x3x12xf32>
 // CHECK: return %[[VAL_2]], %[[VAL_3]] : tensor<1x6x12xf32>, tensor<1x3x12xf32>
 func.func @canonicalize_concat_slice_on_non_concat_axis(%arg0 : tensor<1x12x12xf32>, %arg1 : tensor<1x12x12xf32>) -> (tensor<1x6x12xf32>, tensor<1x3x12xf32>) {
   %0 = "tosa.concat"(%arg0, %arg1) {axis = 2 : i64} : (tensor<1x12x12xf32>, tensor<1x12x12xf32>) -> tensor<1x12x24xf32>
   %1 = "tosa.slice"(%0) {size = array<i64: 1, 6, 12>, start = array<i64: 0, 0, 0>} : (tensor<1x12x24xf32>) -> tensor<1x6x12xf32>
   %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>
+}