[MLIR] Add idempotent trait folding

[ahmedsabie]

Summary

This trait simply adds a fold of f(f(x)) = f(x) when an operation is labelled as idempotent

[ahmedsabie]

Can you push it when you get a chance I don't have push access, thanks!

mlir/include/mlir/IR/OpBase.td

 def AffineScope : NativeOpTrait<"AffineScope">;
 // Op defines an automatic allocation scope.
 def AutomaticAllocationScope : NativeOpTrait<"AutomaticAllocationScope">;
 // Op supports operand broadcast behavior.
 def ResultsBroadcastableShape :
   NativeOpTrait<"ResultsBroadcastableShape">;
 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
+// op op X == op X
+def Idempotent  : NativeOpTrait<"IsIdempotent">;
 // op op X == X
 def Involution  : NativeOpTrait<"IsInvolution">;
 // Op behaves like a constant.
 def ConstantLike : NativeOpTrait<"ConstantLike">;
 // Op behaves like a function.
 def FunctionLike : NativeOpTrait<"FunctionLike">;
 // Op is isolated from above.
 def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;

mlir/include/mlir/IR/OpDefinition.h

 //===----------------------------------------------------------------------===//
 
 namespace OpTrait {
 
 // These functions are out-of-line implementations of the methods in the
 // corresponding trait classes.  This avoids them being template
 // instantiated/duplicated.
 namespace impl {
+OpFoldResult foldIdempotent(Operation *op);
 OpFoldResult foldInvolution(Operation *op);
 LogicalResult verifyZeroOperands(Operation *op);
 LogicalResult verifyOneOperand(Operation *op);
 LogicalResult verifyNOperands(Operation *op, unsigned numOperands);
+LogicalResult verifyIsIdempotent(Operation *op);
 LogicalResult verifyIsInvolution(Operation *op);
 LogicalResult verifyAtLeastNOperands(Operation *op, unsigned numOperands);
 LogicalResult verifyOperandsAreFloatLike(Operation *op);
 LogicalResult verifyOperandsAreSignlessIntegerLike(Operation *op);
 LogicalResult verifySameTypeOperands(Operation *op);
 LogicalResult verifyZeroRegion(Operation *op);
 LogicalResult verifyOneRegion(Operation *op);
 LogicalResult verifyNRegions(Operation *op, unsigned numRegions);

 public:
   static AbstractOperation::OperationProperties getTraitProperties() {
     return static_cast<AbstractOperation::OperationProperties>(
         OperationProperty::Commutative);
   }
 };
 
 /// This class adds property that the operation is an involution.
-/// This means a unary to unary operation "f" that satisfies f(f(x)) = f(x)
+/// This means a unary to unary operation "f" that satisfies f(f(x)) = x
 template <typename ConcreteType>
 class IsInvolution : public TraitBase<ConcreteType, IsInvolution> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     static_assert(ConcreteType::template hasTrait<OneResult>(),
                   "expected operation to produce one result");
     static_assert(ConcreteType::template hasTrait<OneOperand>(),
                   "expected operation to take one operand");
     static_assert(ConcreteType::template hasTrait<SameOperandsAndResultType>(),
                   "expected operation to preserve type");
     // Involution requires the operation to be side effect free as well
     // but currently this check is under a FIXME and is not actually done.
     return impl::verifyIsInvolution(op);
   }
 
   static OpFoldResult foldTrait(Operation *op, ArrayRef<Attribute> operands) {
     return impl::foldInvolution(op);
   }
 };
 
+/// This class adds property that the operation is idempotent.
+/// This means a unary to unary operation "f" that satisfies f(f(x)) = f(x)
+template <typename ConcreteType>
+class IsIdempotent : public TraitBase<ConcreteType, IsIdempotent> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(ConcreteType::template hasTrait<OneResult>(),
+                  "expected operation to produce one result");
+    static_assert(ConcreteType::template hasTrait<OneOperand>(),
+                  "expected operation to take one operand");
+    static_assert(ConcreteType::template hasTrait<SameOperandsAndResultType>(),
+                  "expected operation to preserve type");
+    // Idempotent requires the operation to be side effect free as well
+    // but currently this check is under a FIXME and is not actually done.
+    return impl::verifyIsIdempotent(op);
+  }
+
+  static OpFoldResult foldTrait(Operation *op, ArrayRef<Attribute> operands) {
+    return impl::foldIdempotent(op);
+  }
+};
+
 /// This class verifies that all operands of the specified op have a float type,
 /// a vector thereof, or a tensor thereof.
 template <typename ConcreteType>
 class OperandsAreFloatLike
     : public TraitBase<ConcreteType, OperandsAreFloatLike> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifyOperandsAreFloatLike(op);

mlir/lib/IR/Operation.cpp

 InFlightDiagnostic OpState::emitRemark(const Twine &message) {
   return getOperation()->emitRemark(message);
 }
 
 //===----------------------------------------------------------------------===//
 // Op Trait implementations
 //===----------------------------------------------------------------------===//
 
+OpFoldResult OpTrait::impl::foldIdempotent(Operation *op) {
+  auto *argumentOp = op->getOperand(0).getDefiningOp();
+  if (argumentOp && op->getName() == argumentOp->getName()) {
+    // Replace the outer operation output with the inner operation.
+    return op->getOperand(0);
+  }
+
+  return {};
+}
+
 OpFoldResult OpTrait::impl::foldInvolution(Operation *op) {
   auto *argumentOp = op->getOperand(0).getDefiningOp();
   if (argumentOp && op->getName() == argumentOp->getName()) {
     // Replace the outer involutions output with inner's input.
     return argumentOp->getOperand(0);
   }
 
   return {};
if (auto vec = type.dyn_cast<VectorType>())
     return vec.getElementType();
 
   // Look through tensor<vector<...>> to find the underlying element type.
   if (auto tensor = type.dyn_cast<TensorType>())
     return getTensorOrVectorElementType(tensor.getElementType());
   return type;
 }
 
+LogicalResult OpTrait::impl::verifyIsIdempotent(Operation *op) {
+  // FIXME: Add back check for no side effects on operation.
+  // Currently adding it would cause the shared library build
+  // to fail since there would be a dependency of IR on SideEffectInterfaces
+  // which is cyclical.
+  return success();
+}
+
 LogicalResult OpTrait::impl::verifyIsInvolution(Operation *op) {
   // FIXME: Add back check for no side effects on operation.
   // Currently adding it would cause the shared library build
   // to fail since there would be a dependency of IR on SideEffectInterfaces
   // which is cyclical.
   return success();
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

def TestOpWithVariadicResultsAndFolder: TEST_Op<"op_with_variadic_results_and_folder"> {
   let hasFolder = 1;
 }
 
 def TestCommutativeOp : TEST_Op<"op_commutative", [Commutative]> {
   let arguments = (ins I32:$op1, I32:$op2, I32:$op3, I32:$op4);
   let results = (outs I32);
 }
 
+def TestIdempotentTraitOp
+ : TEST_Op<"op_idempotent_trait",
+           [SameOperandsAndResultType, NoSideEffect, Idempotent]> {
+  let arguments = (ins I32:$op1);
+  let results = (outs I32);
+}
+
 def TestInvolutionTraitNoOperationFolderOp
  : TEST_Op<"op_involution_trait_no_operation_fold",
            [SameOperandsAndResultType, NoSideEffect, Involution]> {
   let arguments = (ins I32:$op1);
   let results = (outs I32);
 }
 
 def TestInvolutionTraitFailingOperationFolderOp

mlir/test/mlir-tblgen/trait.mlir

 // CHECK-SAME:  ([[ARG0:%.+]]: i32)
 func @testInhibitInvolution(%arg0: i32) -> i32 {
   // CHECK: [[OP:%.+]] = "test.op_involution_trait_succesful_operation_fold"([[ARG0]])
   %0 = "test.op_involution_trait_succesful_operation_fold"(%arg0) : (i32) -> i32
   %1 = "test.op_involution_trait_succesful_operation_fold"(%0) : (i32) -> i32
   // CHECK: return [[OP]]
   return %1: i32
 }
+
+//===----------------------------------------------------------------------===//
+// Test that idempotent folding works correctly
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @testSingleIdempotent
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testSingleIdempotent(%arg0 : i32) -> i32 {
+  // CHECK: [[IDEMPOTENT:%.+]] = "test.op_idempotent_trait"([[ARG0]])
+  %0 = "test.op_idempotent_trait"(%arg0) : (i32) -> i32
+  // CHECK: return [[IDEMPOTENT]]
+  return %0: i32
+}
+
+// CHECK-LABEL: func @testDoubleIdempotent
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testDoubleIdempotent(%arg0: i32) -> i32 {
+  // CHECK: [[IDEMPOTENT:%.+]] = "test.op_idempotent_trait"([[ARG0]])
+  %0 = "test.op_idempotent_trait"(%arg0) : (i32) -> i32
+  %1 = "test.op_idempotent_trait"(%0) : (i32) -> i32
+  // CHECK: return [[IDEMPOTENT]]
+  return %1: i32
+}
+
+// CHECK-LABEL: func @testTripleIdempotent
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testTripleIdempotent(%arg0: i32) -> i32 {
+  // CHECK: [[IDEMPOTENT:%.+]] = "test.op_idempotent_trait"([[ARG0]])
+  %0 = "test.op_idempotent_trait"(%arg0) : (i32) -> i32
+  %1 = "test.op_idempotent_trait"(%0) : (i32) -> i32
+  %2 = "test.op_idempotent_trait"(%1) : (i32) -> i32
+  // CHECK: return [[IDEMPOTENT]]
+  return %2: i32
+}