[MLIR] Allow non-binary operations to be commutative

NFC for binary operations.

Differential Revision: https://reviews.llvm.org/D73670

mlir/include/mlir/IR/OperationSupport.h

@@ -56,9 +56,11 @@ class OwningRewritePatternList;
 //===----------------------------------------------------------------------===//
 
 enum class OperationProperty {
-  /// This bit is set for an operation if it is a commutative operation: that
-  /// is a binary operator (two inputs) where "a op b" and "b op a" produce the
-  /// same results.
+  /// This bit is set for an operation if it is a commutative
+  /// operation: that is an operator where order of operands does not
+  /// change the result of the operation.  For example, in a binary
+  /// commutative operation, "a op b" and "b op a" produce the same
+  /// results.
   Commutative = 0x1,
 
   /// This bit is set for operations that have no side effects: that means that

mlir/lib/Transforms/Utils/FoldUtils.cpp

@@ -134,20 +134,19 @@ LogicalResult OperationFolder::tryToFold(
   SmallVector<Attribute, 8> operandConstants;
   SmallVector<OpFoldResult, 8> foldResults;
 
+  // If this is a commutative operation, move constants to be trailing operands.
+  if (op->getNumOperands() >= 2 && op->isCommutative()) {
+    std::stable_partition(
+        op->getOpOperands().begin(), op->getOpOperands().end(),
+        [&](OpOperand &O) { return !matchPattern(O.get(), m_Constant()); });
+  }
+
   // Check to see if any operands to the operation is constant and whether
   // the operation knows how to constant fold itself.
   operandConstants.assign(op->getNumOperands(), Attribute());
   for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i)
     matchPattern(op->getOperand(i), m_Constant(&operandConstants[i]));
 
-  // If this is a commutative binary operation with a constant on the left
-  // side move it to the right side.
-  if (operandConstants.size() == 2 && operandConstants[0] &&
-      !operandConstants[1] && op->isCommutative()) {
-    std::swap(op->getOpOperand(0), op->getOpOperand(1));
-    std::swap(operandConstants[0], operandConstants[1]);
-  }
-
   // Attempt to constant fold the operation.
   if (failed(op->fold(operandConstants, foldResults)))
     return failure();

mlir/test/Transforms/test-canonicalize.mlir

@@ -35,3 +35,19 @@ func @remove_op_with_variadic_results_and_folder(%arg0 : i32, %arg1 : i32) -> (i
   %0, %1 = "test.op_with_variadic_results_and_folder"(%arg0, %arg1) : (i32, i32) -> (i32, i32)
   return %0, %1 : i32, i32
 }
+
+// CHECK-LABEL: func @test_commutative_multi
+// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32, %[[ARG_1:[a-z0-9]*]]: i32)
+func @test_commutative_multi(%arg0: i32, %arg1: i32) -> (i32, i32) {
+  // CHECK: %[[C42:.*]] = constant 42 : i32
+  %c42_i32 = constant 42 : i32
+  // CHECK: %[[C43:.*]] = constant 43 : i32
+  %c43_i32 = constant 43 : i32
+  // CHECK-NEXT: %[[O0:.*]] = "test.op_commutative"(%[[ARG_0]], %[[ARG_1]], %[[C42]], %[[C43]]) : (i32, i32, i32, i32) -> i32
+  %y = "test.op_commutative"(%c42_i32, %arg0, %arg1, %c43_i32) : (i32, i32, i32, i32) -> i32
+
+  // CHECK-NEXT: %[[O1:.*]] = "test.op_commutative"(%[[ARG_0]], %[[ARG_1]], %[[C42]], %[[C43]]) : (i32, i32, i32, i32) -> i32
+  %z = "test.op_commutative"(%arg0, %c42_i32, %c43_i32, %arg1): (i32, i32, i32, i32) -> i32
+  // CHECK-NEXT: return %[[O0]], %[[O1]]
+  return %y, %z: i32, i32
+}

mlir/test/lib/TestDialect/TestOps.td

@@ -639,6 +639,11 @@ def TestOpWithVariadicResultsAndFolder: TEST_Op<"op_with_variadic_results_and_fo
   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);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (Symbol Binding)