[MLIR] Allow non-binary operations to be commutative

[stephenneuendorffer]

Summary

NFC for binary operations.

[merge_guards_bot]

Unit tests: fail. 62314 tests passed, 1 failed and 838 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_class/try_lock.pass.cpp

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/lib/Transforms/Utils/FoldUtils.cpp

142

You should be able to do something like:

return matchPattern(O.get(), m_Constant());

mlir/test/Transforms/test-canonicalize.mlir

48

Don't check SSA values directly:

https://mlir.llvm.org/getting_started/TestingGuide/

[merge_guards_bot]

Unit tests: pass. 62315 tests passed, 0 failed and 838 were skipped.

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: fail. 62513 tests passed, 3 failed and 844 were skipped.

failed: libc++.std/thread/thread_mutex/thread_lock/thread_lock_unique/thread_lock_unique_cons/mutex_time_point.pass.cpp
failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_class/try_lock.pass.cpp
failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_timedmutex_requirements/thread_timedmutex_recursive/try_lock_until.pass.cpp

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[stephenneuendorffer]

Any more comments?

[rriddle]

Can you update the comment in OperationSupport.h:60? It currently mentions binary operations directly.

mlir/lib/Transforms/Utils/FoldUtils.cpp

 /// Tries to perform folding on the given `op`. If successful, populates
 /// `results` with the results of the folding.
 LogicalResult OperationFolder::tryToFold(
     Operation *op, SmallVectorImpl<Value> &results,
     function_ref<void(Operation *)> processGeneratedConstants) {
   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();
 
   // Check to see if the operation was just updated in place.
   if (foldResults.empty())
     return success();
   assert(foldResults.size() == op->getNumResults());

mlir/test/Transforms/test-canonicalize.mlir

 
 // CHECK-LABEL: func @remove_op_with_variadic_results_and_folder
 // CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32, %[[ARG_1:[a-z0-9]*]]: i32)
 func @remove_op_with_variadic_results_and_folder(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
   // CHECK-NEXT: return %[[ARG_0]], %[[ARG_1]]
   %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

 }
 
 def TestOpWithVariadicResultsAndFolder: TEST_Op<"op_with_variadic_results_and_folder"> {
   let arguments = (ins Variadic<I32>:$operands);
   let results = (outs Variadic<I32>);
   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)
 
 // Test symbol binding.
 def OpSymbolBindingA : TEST_Op<"symbol_binding_a", []> {
   let arguments = (ins I32:$operand, I64Attr:$attr);
   let results = (outs I32);
 }