[PATCH] D73670: [MLIR] Allow non-binary operations to be commutative

[Stephen Neuendorffer via Phabricator via llvm-commits]

stephenneuendorffer updated this revision to Diff 241320.

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D73670/new/

https://reviews.llvm.org/D73670

Files:
  mlir/lib/Transforms/Utils/FoldUtils.cpp
  mlir/test/Transforms/test-canonicalize.mlir
  mlir/test/lib/TestDialect/TestOps.td


Index: mlir/test/lib/TestDialect/TestOps.td
===================================================================
--- mlir/test/lib/TestDialect/TestOps.td
+++ mlir/test/lib/TestDialect/TestOps.td
@@ -639,6 +639,11 @@
   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)
 
Index: mlir/test/Transforms/test-canonicalize.mlir
===================================================================
--- mlir/test/Transforms/test-canonicalize.mlir
+++ mlir/test/Transforms/test-canonicalize.mlir
@@ -35,3 +35,18 @@
   %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
+func @test_commutative_multi(%arg0: i32, %arg1: i32) -> (i32, i32) {
+  // CHECK: %c42_i32 = constant 42 : i32
+  %c42_i32 = constant 42 : i32
+  // CHECK: %c43_i32 = constant 43 : i32
+  %c43_i32 = constant 43 : i32
+  // CHECK-NEXT: %0 = "test.op_commutative"(%arg0, %arg1, %c42_i32, %c43_i32) : (i32, i32, i32, i32) -> i32
+  %y = "test.op_commutative"(%c42_i32, %arg0, %arg1, %c43_i32) : (i32, i32, i32, i32) -> i32
+
+  // CHECK-NEXT: %1 = "test.op_commutative"(%arg0, %arg1, %c42_i32, %c43_i32) : (i32, i32, i32, i32) -> i32
+  %z = "test.op_commutative"(%arg0, %c42_i32, %c43_i32, %arg1): (i32, i32, i32, i32) -> i32
+  // CHECK-NEXT: return %0, %1
+  return %y, %z: i32, i32
+}
Index: mlir/lib/Transforms/Utils/FoldUtils.cpp
===================================================================
--- mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -134,20 +134,22 @@
   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) {
+                            Attribute c;
+                            matchPattern(O.get(), m_Constant(&c));
+                            return (bool)!c;
+                          });
+  }
+
   // 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();


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D73670.241320.patch
Type: text/x-patch
Size: 3217 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200130/721908f0/attachment.bin>