[mlir] Use strictly less than comparison in Commutative operand sorting

[amandatang]

Summary

There is an assertion failure within stable_sort on the comparison function that is passed. stable_sort expects a strictly less than comparison function, which means it should return false for equivalent values.

Note that equivalent operands have the same key size and key values at each index. When the algorithm reaches a keyIndex that exceeds the operands' key sizes, it checks commOperandA first and separately from commOperandB. If the current keyIndex exceeds commOperandA's key size and commOperandA's ancestorQueue is empty, then the comparison function returns true, implying commOperandA is smaller regardless of commOperandB's key.

To fix this, before returning true if commOperandA's ancestorQueue is empty, commOperandB's key size and ancestorQueue is checked for equivalence.

This patch also includes an improvement that uses the OpTraitRewritePattern which is a wrapper for the RewritePattern class that only performs matchAndRewrite on Ops with a specified trait.

Finally, Example 2 in the comments is corrected. The key associated with operand %2 is missing a block argument entry at the end. Additionally, the final sorted operand ordering should be %4, %2, %3, %1.

mlir/lib/Transforms/Utils/CommutativityUtils.cpp

 /// Here,
 /// 1. The key associated with %1 is:
 ///     `{
 ///       {CONSTANT_OP, "foo.const"}
 ///      }`
 /// 2. The key associated with %2 is:
 ///     `{
 ///       {NON_CONSTANT_OP, "foo.mul"},
+///       {BLOCK_ARGUMENT, ""},
 ///       {BLOCK_ARGUMENT, ""}
 ///      }`
 /// 3. The key associated with %3 is:
 ///     `{
 ///       {NON_CONSTANT_OP, "foo.mul"},
 ///       {NON_CONSTANT_OP, "foo.mul"},
 ///       {CONSTANT_OP, "foo.const"},
 ///       {BLOCK_ARGUMENT, ""},
 ///       {BLOCK_ARGUMENT, ""}
 ///      }`
 /// 4. The key associated with %4 is:
 ///     `{
 ///       {NON_CONSTANT_OP, "foo.add"},
 ///       {NON_CONSTANT_OP, "foo.mul"},
 ///       {CONSTANT_OP, "foo.const"},
 ///       {BLOCK_ARGUMENT, ""},
 ///       {BLOCK_ARGUMENT, ""}
 ///      }`
 ///
 /// Thus, the sorted `foo.commutative` is:
-/// %5 = foo.commutative %4, %3, %2, %1
-class SortCommutativeOperands : public RewritePattern {
-public:
+/// %5 = foo.commutative %4, %2, %3, %1
+struct SortCommutativeOperands final
+    : public OpTraitRewritePattern<OpTrait::IsCommutative> {
   SortCommutativeOperands(MLIRContext *context)
-      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/5, context) {}
+      : OpTraitRewritePattern<OpTrait::IsCommutative>(context, /*benefit=*/5) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     // Custom comparator for two commutative operands, which returns true iff
     // the "key" of `constCommOperandA` < the "key" of `constCommOperandB`,
     // i.e.,
     // 1. In the first unequal pair of corresponding AncestorKeys, the
     // AncestorKey in `constCommOperandA` is smaller, or,
     // 2. Both the AncestorKeys in every pair are the same and the size of
auto commutativeOperandComparator =
               const_cast<std::unique_ptr<CommutativeOperand> &>(
                   constCommOperandB);
 
           // Iteratively perform the BFS's of both operands until an order among
           // them can be determined.
           unsigned keyIndex = 0;
           while (true) {
             if (commOperandA->key.size() <= keyIndex) {
+              // Comparator must return false for equal elements
+              // B is only larger if its key size is larger than the current
+              // index or its ancestor queue is not empty
               if (commOperandA->ancestorQueue.empty())
-                return true;
+                return commOperandB->key.size() > keyIndex ||
+                       !commOperandB->ancestorQueue.empty();
               commOperandA->popFrontAndPushAdjacentUnvisitedAncestors();
               commOperandA->refreshKey();
             }
             if (commOperandB->key.size() <= keyIndex) {
               if (commOperandB->ancestorQueue.empty())
                 return false;
               commOperandB->popFrontAndPushAdjacentUnvisitedAncestors();
               commOperandB->refreshKey();
             }
             if (commOperandA->ancestorQueue.empty() ||
                 commOperandB->ancestorQueue.empty())
               return commOperandA->key.size() < commOperandB->key.size();
             if (commOperandA->key[keyIndex] < commOperandB->key[keyIndex])
               return true;
             if (commOperandB->key[keyIndex] < commOperandA->key[keyIndex])
               return false;
             keyIndex++;
           }
         };
 
-    // If `op` is not commutative, do nothing.
-    if (!op->hasTrait<OpTrait::IsCommutative>())
-      return failure();
-
     // Populate the list of commutative operands.
     SmallVector<Value, 2> operands = op->getOperands();
     SmallVector<std::unique_ptr<CommutativeOperand>, 2> commOperands;
     for (Value operand : operands) {
       std::unique_ptr<CommutativeOperand> commOperand =
           std::make_unique<CommutativeOperand>();
       commOperand->operand = operand;
       commOperand->pushAncestor(operand.getDefiningOp());

mlir/test/Transforms/test-commutativity-utils.mlir

func.func @test_large_pattern(%arg0 : i32, %arg1 : i32) -> i32 {
   %21 = arith.divsi %17, %20 : i32
 
   // CHECK-NEXT: %[[RESULT:.*]] = "test.op_large_commutative"(%[[VAL16]], %[[VAL19]], %[[VAL19]], %[[VAL21]], %[[VAL6]], %[[VAL11]], %[[VAL15]])
   %result = "test.op_large_commutative"(%16, %6, %11, %15, %19, %21, %19): (i32, i32, i32, i32, i32, i32, i32) -> i32
 
   // CHECK-NEXT: return %[[RESULT]]
   return %result : i32
 }
+
+// CHECK-LABEL: @test_equal_ancestor_trees
+func.func @test_equal_ancestor_trees(%arg0 : i32, %arg1: i32) -> i32 {
+  // CHECK-NEXT: arith.addi
+  %0 = arith.addi %arg1, %arg1 : i32
+
+  // CHECK-NEXT: arith.addi
+  %1 = arith.addi %arg0, %arg0 : i32
+
+  // CHECK: %[[ARITH_MULI1:.*]] = arith.muli
+  %2 = arith.muli %0, %0 : i32
+
+  // CHECK-NEXT: %[[ARITH_MULI2:.*]] = arith.muli
+  %3 = arith.muli %1, %1 : i32
+
+  // Without additional logic to differentiate between block arguments, %2 and %3 are treated
+  // as equivalent, so the sorting logic keeps the original order
+  // CHECK-NEXT: %[[RESULT:.*]] = "test.op_commutative2"(%[[ARITH_MULI1]], %[[ARITH_MULI2]])
+  %result = "test.op_commutative2"(%2, %3): (i32, i32) -> i32
+
+  // CHECK-NEXT: return %[[RESULT]]
+  return %result : i32
+}