[mlir][tosa] Fold consecutive tosa.abs

[Lewuathe]

Summary

Consecutive tosa.abs can be fold as single abs operation since the second one has no impact.

[mgehre-amd]

Looks good!

In general, MLIR could have an Idempotence trait on ops and then a general facility in the folder.
In TOSA specifically, I think also floor, ceil and clamp (with same min/max) are idempotent.

[jpienaar]

In D150836#4355445, @mgehre-amd wrote:

Looks good!

In general, MLIR could have an Idempotence trait on ops and then a general facility in the folder.
In TOSA specifically, I think also floor, ceil and clamp (with same min/max) are idempotent.

Funny you should mention that

https://mlir.llvm.org/doxygen/classmlir_1_1OpTrait_1_1IsIdempotent.html

[Lewuathe]

In D150836#4355445, @mgehre-amd wrote:

Looks good!

In general, MLIR could have an Idempotence trait on ops and then a general facility in the folder.
In TOSA specifically, I think also floor, ceil and clamp (with same min/max) are idempotent.

@mgehre-amd @jpienarr Okay. I’ll add the IsIdempotent trait for these ops in a separate change. Thanks!

[Lewuathe]

In TOSA specifically, I think also floor, ceil and clamp (with same min/max) are idempotent.

Idempotent trait seems to expect the op not to change the type itself. They cannot be applied to these ops as it is.

static_assert failed due to requirement 'CeilOp::hasTrait()' "expected operation to preserve type"

@mgehre-amd

[mgehre-amd]

In D150836#4366820, @Lewuathe wrote:

In TOSA specifically, I think also floor, ceil and clamp (with same min/max) are idempotent.

Idempotent trait seems to expect the op not to change the type itself. They cannot be applied to these ops as it is.

static_assert failed due to requirement 'CeilOp::hasTrait()' "expected operation to preserve type"

@mgehre-amd

I don't understand what you mean. How does floor for example change the type?

[eric-k256]

I think the issue is that floor/ceil/clamp don't have the SameOperandsAndResultType trait which is required by the Idempotent trait. It should be okay to add for floor/ceil/clamp as they shouldn't change type.

We can't add SameOperandsAndResultType to the TOSA elementwise binary operators because of the implicit broadcast in TOSA. In those cases the difference in shape between the input and output would cause failures. (Not a factor here, but noting it for reference)

We can't put Idempotent on clamp because the Idempotent fold operation doesn't check the attributes, so could potentially fold improperly. It likely needs an op specific folder.

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

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

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

OpFoldResult tosa::NegateOp::fold(FoldAdaptor adaptor) {
   auto input = getInput1();
   // Element-wise negate(negate(x)) = x
   if (auto op = input.getDefiningOp<tosa::NegateOp>()) {
     return op.getInput1();
   }
 
   return {};
 }
+
+OpFoldResult tosa::AbsOp::fold(FoldAdaptor adaptor) {
+  auto input = getInput1();
+  // Element-wise abs(abs(x)) = abs(x)
+  if (auto op = input.getDefiningOp<tosa::AbsOp>()) {
+    return input;
+  }
+
+  return {};
+}

mlir/test/Dialect/Tosa/canonicalize.mlir

 // CHECK-LABEL: @fold_negate_negate
 func.func @fold_negate_negate(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
   // CHECK: return %arg{{.*}} : tensor<?x1xf32>
   %0 = "tosa.negate"(%arg0) : (tensor<?x1xf32>) -> tensor<?x1xf32>
   %1 = "tosa.negate"(%0) : (tensor<?x1xf32>) -> tensor<?x1xf32>
   return %1 : tensor<?x1xf32>
 }
 
+// -----
+
+// CHECK-LABEL: @fold_abs_abs
+func.func @fold_abs_abs(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
+  // CHECK: %[[ABS:.*]] = "tosa.abs"(%arg{{.*}}) : (tensor<?x1xf32>) -> tensor<?x1xf32>
+  // CHECK: return %[[ABS]] : tensor<?x1xf32>
+  %0 = "tosa.abs"(%arg0) : (tensor<?x1xf32>) -> tensor<?x1xf32>
+  %1 = "tosa.abs"(%0) : (tensor<?x1xf32>) -> tensor<?x1xf32>
+  return %1 : tensor<?x1xf32>
+}