[MLIR] Allow `Idempotent` trait to be applied to binary ops.

[chr1sj0nes]

Summary

Add Idempotent trait to arith.{andi,ori}.

[chr1sj0nes]

Formatting fix.

[mehdi_amini]

mlir/lib/IR/Operation.cpp

638

Nit: no need for else after return

[chr1sj0nes]

Thanks for the review Mehdi! Can you please submit if you are happy, as I don't have permissions.

mlir/lib/IR/Operation.cpp

638

I think the else is necessary here, as the if branch doesn't always return. Without the else, it might call op->getOperand(1) on an op with only one operand.

[lattner]

Cool. Is there any reason that idompotence be limited to unary and binary ops? Variadic "or" operations like comb.or in CIRCT seem like they should be able to participate as well.

[chr1sj0nes]

In D114574#3156930, @lattner wrote:

Cool. Is there any reason that idompotence be limited to unary and binary ops? Variadic "or" operations like comb.or in CIRCT seem like they should be able to participate as well.

I guess the question for me there would be, is that stretching the standard terminology too much? The Wikipedia article (https://en.wikipedia.org/wiki/Idempotence), for example, makes no mention of variadic ops that I can see. Perhaps a question for the maths / compsci theorists.

[jpienaar]

In D114574#3160947, @chr1sj0nes wrote:

In D114574#3156930, @lattner wrote:

Cool. Is there any reason that idompotence be limited to unary and binary ops? Variadic "or" operations like comb.or in CIRCT seem like they should be able to participate as well.

I guess the question for me there would be, is that stretching the standard terminology too much? The Wikipedia article (https://en.wikipedia.org/wiki/Idempotence), for example, makes no mention of variadic ops that I can see. Perhaps a question for the maths / compsci theorists.

I think technically it is only defined on unary operations, for the binary case we require both operands be the same, so the function that is idempotent is rather the unary function with value x that invokes the binary op with x,x and in that sense variadic is merely a unary function with a vector of values (but being a repetition of 1 input) as input. So the same generalization seem to apply, now it is stretching the term slightly, but I can't think of one that is less confusing while accurate.

[lattner]

We can argue about specific definitions, but (for example) the wikipedia definition doesn't discuss arity at all: " is the property of certain operations in mathematics and computer science whereby they can be applied multiple times without changing the result beyond the initial application." If you're going to define this for binary, then you need to define what it means for multiple operands. Given that, I don't see how 2 is more special than N :-)

Going beyond theory, I think we should also be guided by practical considerations - is it useful for MLIR-based compilers to have one definition or the other?

[mehdi_amini]

The generalization to variadic makes sense to me, in particular for an op like or which seems like it can be expressed like a compressed form of a chain of the binary op as long as it is associative : or(x, y, z) is equivalent to or(or(x, y), z) or or(x, or(y, z)).

Asked differently: I'm not sure what we get from limiting the Idempotent trait to unary and binary ops?

[rriddle]

In D114574#3163055, @mehdi_amini wrote:

The generalization to variadic makes sense to me, in particular for an op like or which seems like it can be expressed like a compressed form of a chain of the binary op as long as it is associative : or(x, y, z) is equivalent to or(or(x, y), z) or or(x, or(y, z)).

Asked differently: I'm not sure what we get from limiting the Idempotent trait to unary and binary ops?

+1. As long as our definition in the face of variadic is clear, variadic operations that don't conform to that definition could just not add the trait (which would be the same if we limited it to unary/binary anyways).

[mehdi_amini]

In D114574#3162232, @lattner wrote:

We can argue about specific definitions, but (for example) the wikipedia definition doesn't discuss arity at all

To be fair: while the intro gives a fairly generic definition, if you continue in the "definition" section it actually only talks about binary and not unary at all I believe. The examples are full of mostly binary operators as well, until you get to the "function" section where unary examples are provided again.

I don't think we should anchor to closely to wikipedia here, the article isn't entirely consistent and looks quite informal and context dependent as far as I can tell.

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td

def Arith_RemSIOp : Arith_IntBinaryOp<"remsi"> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AndIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_AndIOp : Arith_IntBinaryOp<"andi", [Commutative]> {
+def Arith_AndIOp : Arith_IntBinaryOp<"andi", [Commutative, Idempotent]> {
   let summary = "integer binary and";
   let description = [{
     The `andi` operation takes two operands and returns one result, each of
     these is required to be the same type. This type may be an integer scalar
     type, a vector whose element type is integer, or a tensor of integers. It
     has no standard attributes.
 
     Example:
def Arith_AndIOp : Arith_IntBinaryOp<"andi", [Commutative, Idempotent]> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // OrIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_OrIOp : Arith_IntBinaryOp<"ori", [Commutative]> {
+def Arith_OrIOp : Arith_IntBinaryOp<"ori", [Commutative, Idempotent]> {
   let summary = "integer binary or";
   let description = [{
     The `ori` operation takes two operands and returns one result, each of these
     is required to be the same type. This type may be an integer scalar type, a
     vector whose element type is integer, or a tensor of integers. It has no
     standard attributes.
 
     Example:

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
+// op op X == op X (unary) / X op X == X (binary)
 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.

mlir/include/mlir/IR/OpDefinition.h

public:
   }
 
   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)
+/// This means a unary to unary operation "f" that satisfies f(f(x)) = f(x),
+/// or a binary operation "g" that satisfies g(x, x) = 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<OneOperand>() ||
+                      ConcreteType::template hasTrait<NOperands<2>::Impl>(),
+                  "expected operation to take one or two operands");
     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) {

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

 OpFoldResult arith::AndIOp::fold(ArrayRef<Attribute> operands) {
   /// and(x, 0) -> 0
   if (matchPattern(getRhs(), m_Zero()))
     return getRhs();
   /// and(x, allOnes) -> x
   APInt intValue;
   if (matchPattern(getRhs(), m_ConstantInt(&intValue)) && intValue.isAllOnes())
     return getLhs();
-  /// and(x, x) -> x
-  if (getLhs() == getRhs())
-    return getRhs();
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a & b; });
 }
 
 //===----------------------------------------------------------------------===//
 // OrIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::OrIOp::fold(ArrayRef<Attribute> operands) {
   /// or(x, 0) -> x
   if (matchPattern(getRhs(), m_Zero()))
     return getLhs();
-  /// or(x, x) -> x
-  if (getLhs() == getRhs())
-    return getRhs();
   /// or(x, <all ones>) -> <all ones>
   if (auto rhsAttr = operands[1].dyn_cast_or_null<IntegerAttr>())
     if (rhsAttr.getValue().isAllOnes())
       return rhsAttr;
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a | b; });
 }

mlir/lib/IR/Operation.cpp

InFlightDiagnostic OpState::emitRemark(const Twine &message) {
   return getOperation()->emitRemark(message);
 }
 
 //===----------------------------------------------------------------------===//
 // Op Trait implementations
 //===----------------------------------------------------------------------===//
 
 OpFoldResult OpTrait::impl::foldIdempotent(Operation *op) {
+  if (op->getNumOperands() == 1) {
     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);
     }
+  } else if (op->getOperand(0) == op->getOperand(1)) {
+    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.

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

 
 def TestIdempotentTraitOp
  : TEST_Op<"op_idempotent_trait",
            [SameOperandsAndResultType, NoSideEffect, Idempotent]> {
   let arguments = (ins I32:$op1);
   let results = (outs I32);
 }
 
+def TestIdempotentTraitBinaryOp
+    : TEST_Op<"op_idempotent_trait_binary",
+              [SameOperandsAndResultType, NoSideEffect, Idempotent]> {
+  let arguments = (ins I32:$op1, I32:$op2);
+  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

 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
 }
+
+// CHECK-LABEL: func @testBinaryIdempotent
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testBinaryIdempotent(%arg0 : i32) -> i32 {
+  %0 = "test.op_idempotent_trait_binary"(%arg0, %arg0) : (i32, i32) -> i32
+  // CHECK: return [[ARG0]]
+  return %0: i32
+}