[mlir][CSE] Add ability to remove commutative operations

[clementval]

Summary

This patch takes advantage of the Commutative trait on operation
to remove identical commutative operations where the operands are swapped.

The second operation below can be removed since arith.addi is commutative.

%1 = arith.addi %a, %b : i32
%2 = arith.addi %b, %a : i32

[clementval]

• Rebase
• Fix check line
• Add logic for isEquivalentTo

[rriddle]

mlir/lib/IR/OperationSupport.cpp

647

nit: Drop the llvm:: here, and can you pre-reserve the number of operands?

648

Can you just use a range loop here?

652

This should be using hashOperands to hash operands of the operation. Can you just use an optional storage idiom (like suggested below) and reuse the existing code path?

752–753

The commutative support here isn't comparing the same as the non-commutative path. Can you refactor to use optional storage? E.g. using an idiom like:

ValueRange lhsOperands = lhs->getOperands(), rhsOperands = rhs->getOperands();
SmallVector<Value> lhsOperandStorage, rhsOperandStorage;
if (lhs->hasTrait<mlir::OpTrait::IsCommutative>()) {
  lhsOperandStrorage.append(lhsOperands.begin(), lhsOperands.end());
  llvm::sort(lhsOperandStrorage, ...);
  lhsOperands = lhsOperandStrorage;

  ...
}

// Check mapping of operands and results.
if (!checkValueRangeMapping(lhsOperands, rhsOperands, mapOperands))

[mehdi_amini]

Nice! LG with River's comment.

[clementval]

Address review comments

[clementval]

@rriddle ping

[rriddle]

LGTM, nice!

mlir/lib/IR/OperationSupport.cpp

731

nit: Could you add a newline before this? to separate the lhs initialization from the rhs?

[mehdi_amini]

mlir/lib/IR/OperationSupport.cpp

737

This has a bug: when OperationEquivalence is used with mapOperands, for comparing different blocks/functions for example, sorting based on the pointer values isn't correct: the point values are different in the two blocks and the normalization done here isn't stable.

[mehdi_amini]

mlir/lib/IR/OperationSupport.cpp

737

Could we try instead to move this to a canonicalization: we could sort operands but we'd need a stable sort based on some numbering in the region maybe.

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

Will block/function have the IsCommutative trait? I'm not sure I follow you here. Do you have an example maybe?

[mehdi_amini]

mlir/lib/IR/OperationSupport.cpp

737

func @foo1(%a : i32, %b : i32) -> i32 {
  %1 = arith.addi %a, %b : i32
  return %1 : i32
}
func @foo2(%a : i32, %b : i32) -> i32 {
  %1 = arith.addi %a, %b : i32
  return %1 : i32
}

We should be able to use OperationEquivalence on @foo1 and @foo2. However the pointer values of %a and %b will obviously be different and can sort differently right?

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

Ok it makes more sense. Maybe we can have a special OperationEquivalence for CSE that does not applies in other cases.

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

@mehdi_amini ping. I was away for parental leave for the last 5 weeks. Happy to work on this if we have a solution.

[mehdi_amini]

mlir/lib/IR/OperationSupport.cpp

737

I don't know what's the best solution, but the current bug has been quite an issue for some internal customers of us. If you can propose something here that would be great!

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

I had another look at this now that I'm done with upstreaming. The FuncOp is not commutative so your example the OperationEquivalence will not make use of the pointer and fall back to the initial algo.

[mehdi_amini]

mlir/lib/IR/OperationSupport.cpp

737

I don't quite get what you mean by "the initial algo" here?

Here is some sample of the code we have internally (I didn't write this, but it's using operation equivalence):

// Implements func op equivalence. Func ops that are equivalent will return the
// same hash value for `getHashValue()` and true for `isEqual()`.
struct FuncInfo : public llvm::DenseMapInfo<mlir::func::FuncOp> {
  static unsigned getHashValue(mlir::func::FuncOp func) {
    llvm::hash_code hash{};
    for (mlir::Operation& op : func.getOps()) {
      // No need to look at the regions owned by the op, if any, because the
      // current `isEqual()` implementation doesn't support ops with regions.
      hash = llvm::hash_combine(
          hash, OperationEquivalence::computeHash(
                    &op,
                    /*hashOperands=*/OperationEquivalence::ignoreHashValue,
                    /*hashResults=*/OperationEquivalence::ignoreHashValue,
                    OperationEquivalence::IgnoreLocations));
    }
    return hash;
  }

  static bool isEqual(mlir::func::FuncOp lhs, mlir::func::FuncOp rhs) {
    // Exact match for special keys.
    if (lhs == getEmptyKey() || rhs == getEmptyKey() ||
        lhs == getTombstoneKey() || rhs == getTombstoneKey()) {
      return lhs == rhs;
    }

    if (lhs.getFunctionType() != rhs.getFunctionType()) {
      return false;
    }

    // Compare func ops' attributes excluding their names.
    {
      mlir::NamedAttrList lhs_attrs(lhs->getAttrDictionary());
      lhs_attrs.erase(lhs.getSymNameAttrName());
      mlir::NamedAttrList rhs_attrs(rhs->getAttrDictionary());
      rhs_attrs.erase(rhs.getSymNameAttrName());

      if (lhs_attrs != rhs_attrs) {
        return false;
      }
    }

    // Only support single-block func ops for now.
    if (!lhs.getBody().hasOneBlock() || !rhs.getBody().hasOneBlock()) {
      return false;
    }
    if (lhs.front().getOperations().size() !=
        rhs.front().getOperations().size()) {
      return false;
    }

    // Mapping from rhs' values to lhs' values.
    llvm::DenseMap<mlir::Value, mlir::Value> rhs_to_lhs_values;
    for (const auto& [lhs_arg, rhs_arg] :
         llvm::zip(lhs.getArguments(), rhs.getArguments())) {
      rhs_to_lhs_values[rhs_arg] = lhs_arg;
    }

    // Compare individual op equvialence with `mlir::OperationEquivalence`.
    // The current implementation only supports ops without any regions.
    for (const auto& [lhs_op, rhs_op] : llvm::zip(lhs.front(), rhs.front())) {
      if (lhs_op.getNumRegions() != 0 || rhs_op.getNumRegions() != 0) {
        return false;
      }

      auto map_operands = [&](mlir::Value lhs_value, mlir::Value rhs_value) {
        const auto it = rhs_to_lhs_values.find(rhs_value);
        return mlir::success(lhs_value == it->second);
      };
      if (!OperationEquivalence::isEquivalentTo(
              &lhs_op, &rhs_op, /*mapOperands=*/map_operands,
              /*mapResults=*/OperationEquivalence::ignoreValueEquivalence,
              OperationEquivalence::IgnoreLocations)) {
        return false;
      }

      for (const auto& [lhs_result, rhs_result] :
           llvm::zip(lhs_op.getResults(), rhs_op.getResults())) {
        rhs_to_lhs_values[rhs_result] = lhs_result;
      }
    }

    return true;
  }
};

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

I don't quite get what you mean by "the initial algo" here?

What I meant is that if the operation your are comparing do not have the commutative trait then the comparison is actually the one done before this patch. This patch only change the behavior for commutative operation and should not impact your code. Let me check with your code and the small example your provide in a comment above.

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

I just tested your code on your example:

func.func @foo1(%a : i32, %b : i32) -> i32 {
  %1 = arith.addi %a, %b : i32
  return %1 : i32
}
func.func @foo2(%a : i32, %b : i32) -> i32 {
  %1 = arith.addi %a, %b : i32
  return %1 : i32
}

The two functions are equal and their hashcode are identical. The problem might come from something else.

[mehdi_amini]

mlir/lib/IR/OperationSupport.cpp

737

Isn’t this a matter of « chance »? That is it depends on the Pointer values which may or may not trigger from a textual example?

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

The pointer are not involved here. FuncOp does not have the commutative trait.

[mehdi_amini]

mlir/lib/IR/OperationSupport.cpp

737

But we will call operation equivalence on the body and each op in the body. The add ops are commutative, so there I would expect your normalization to be involved and the add may or may not be equivalent depending on the order of the pointers, wouldn’t it?

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

Ok now that makes sense to me. I'll have a second look.

[clementval]

mlir/lib/IR/OperationSupport.cpp

737

A simple solution is to restore the previous behavior and makes the pointers equivalence for commutative operations configurable through a flag. An attempt at this is available here: https://reviews.llvm.org/D129480

mlir/lib/IR/OperationSupport.cpp

llvm::hash_code OperationEquivalence::computeHash(
   // Hash operations based upon their:
   //   - Operation Name
   //   - Attributes
   //   - Result Types
   llvm::hash_code hash = llvm::hash_combine(
       op->getName(), op->getAttrDictionary(), op->getResultTypes());
 
   //   - Operands
-  for (Value operand : op->getOperands())
+  ValueRange operands = op->getOperands();
+  SmallVector<Value> operandStorage;
+  if (op->hasTrait<mlir::OpTrait::IsCommutative>()) {
+    operandStorage.append(operands.begin(), operands.end());
+    llvm::sort(operandStorage, [](Value a, Value b) -> bool {
+      return a.getAsOpaquePointer() < b.getAsOpaquePointer();
+    });
+    operands = operandStorage;
+  }
+  for (Value operand : operands)
     hash = llvm::hash_combine(hash, hashOperands(operand));
+
   //   - Operands
   for (Value result : op->getResults())
     hash = llvm::hash_combine(hash, hashResults(result));
   return hash;
 }
 
 static bool
 isRegionEquivalentTo(Region *lhs, Region *rhs,
                      function_ref<LogicalResult(Value, Value)> mapOperands,
                      function_ref<LogicalResult(Value, Value)> mapResults,
                      OperationEquivalence::Flags flags) {
   DenseMap<Block *, Block *> blocksMap;
   auto blocksEquivalent = [&](Block &lBlock, Block &rBlock) {
if (lhs->getName() != rhs->getName() ||
       lhs->getNumRegions() != rhs->getNumRegions() ||
       lhs->getNumSuccessors() != rhs->getNumSuccessors() ||
       lhs->getNumOperands() != rhs->getNumOperands() ||
       lhs->getNumResults() != rhs->getNumResults())
     return false;
   if (!(flags & IgnoreLocations) && lhs->getLoc() != rhs->getLoc())
     return false;
 
+  ValueRange lhsOperands = lhs->getOperands(), rhsOperands = rhs->getOperands();
+  SmallVector<Value> lhsOperandStorage, rhsOperandStorage;
+  if (lhs->hasTrait<mlir::OpTrait::IsCommutative>()) {
+    lhsOperandStorage.append(lhsOperands.begin(), lhsOperands.end());
+    llvm::sort(lhsOperandStorage, [](Value a, Value b) -> bool {
+      return a.getAsOpaquePointer() < b.getAsOpaquePointer();
+    });
+    lhsOperands = lhsOperandStorage;
+    rhsOperandStorage.append(rhsOperands.begin(), rhsOperands.end());
+    llvm::sort(rhsOperandStorage, [](Value a, Value b) -> bool {
+      return a.getAsOpaquePointer() < b.getAsOpaquePointer();
+    });
+    rhsOperands = rhsOperandStorage;
+  }
   auto checkValueRangeMapping =
       [](ValueRange lhs, ValueRange rhs,
          function_ref<LogicalResult(Value, Value)> mapValues) {
         for (auto operandPair : llvm::zip(lhs, rhs)) {
           Value curArg = std::get<0>(operandPair);
           Value otherArg = std::get<1>(operandPair);
           if (curArg.getType() != otherArg.getType())
             return false;
           if (failed(mapValues(curArg, otherArg)))
             return false;
         }
         return true;
       };
   // Check mapping of operands and results.
-  if (!checkValueRangeMapping(lhs->getOperands(), rhs->getOperands(),
-                              mapOperands))
+  if (!checkValueRangeMapping(lhsOperands, rhsOperands, mapOperands))
     return false;
   if (!checkValueRangeMapping(lhs->getResults(), rhs->getResults(), mapResults))
     return false;
   for (auto regionPair : llvm::zip(lhs->getRegions(), rhs->getRegions()))
     if (!isRegionEquivalentTo(&std::get<0>(regionPair),
                               &std::get<1>(regionPair), mapOperands, mapResults,
                               flags))
       return false;
   return true;
 }

mlir/test/Transforms/cse.mlir

func @dont_remove_duplicated_read_op_with_sideeffecting() -> i32 {
   %0 = "test.op_with_memread"() : () -> (i32)
   "test.op_with_memwrite"() : () -> ()
   // CHECK: %[[READ_VALUE1:.*]] = "test.op_with_memread"() : () -> i32
   %1 = "test.op_with_memread"() : () -> (i32)
   // CHECK-NEXT: %{{.*}} = arith.addi %[[READ_VALUE0]], %[[READ_VALUE1]] : i32
   %2 = arith.addi %0, %1 : i32
   return %2 : i32
 }
+
+/// This test is checking that identical commutative operation are gracefully
+/// handled but the CSE pass.
+// CHECK-LABEL: func @check_cummutative_cse
+func @check_cummutative_cse(%a : i32, %b : i32) -> i32 {
+  // CHECK: %[[ADD1:.*]] = arith.addi %{{.*}}, %{{.*}} : i32
+  %1 = arith.addi %a, %b : i32
+  %2 = arith.addi %b, %a : i32
+  // CHECK-NEXT:  arith.muli %[[ADD1]], %[[ADD1]] : i32
+  %3 = arith.muli %1, %2 : i32
+  return %3 : i32
+}