Page MenuHomePhabricator

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

Authored by clementval on Apr 11 2022, 3:28 AM.

Details

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

Diff Detail

Event Timeline

clementval created this revision.Apr 11 2022, 3:28 AM
Herald added a project: Restricted Project. · View Herald TranscriptApr 11 2022, 3:28 AM
clementval requested review of this revision.Apr 11 2022, 3:28 AM
  • Rebase
  • Fix check line
  • Add logic for isEquivalentTo
rriddle added inline comments.Apr 11 2022, 10:22 AM
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?

753–754

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))

Nice! LG with River's comment.

clementval marked 3 inline comments as done.

Address review comments

clementval marked an inline comment as done.Apr 12 2022, 12:57 AM
rriddle accepted this revision.Apr 14 2022, 6:59 PM

LGTM, nice!

mlir/lib/IR/OperationSupport.cpp
731

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

This revision is now accepted and ready to land.Apr 14 2022, 6:59 PM
mehdi_amini added inline comments.Apr 20 2022, 5:36 PM
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 added inline comments.Apr 20 2022, 6:47 PM
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 added inline comments.Apr 21 2022, 12:47 AM
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 added inline comments.Apr 21 2022, 9:18 AM
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 added inline comments.Apr 21 2022, 9:29 AM
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 added inline comments.Jun 7 2022, 4:36 AM
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 added inline comments.Jun 7 2022, 7:04 PM
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 added inline comments.Jul 7 2022, 12:44 PM
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 added inline comments.Jul 7 2022, 1:50 PM
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 added inline comments.Jul 8 2022, 12:42 AM
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 added inline comments.Jul 8 2022, 2:06 AM
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 added inline comments.Jul 8 2022, 2:14 AM
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 added inline comments.Jul 8 2022, 2:16 AM
mlir/lib/IR/OperationSupport.cpp
737

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

mehdi_amini added inline comments.Jul 8 2022, 11:44 AM
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 added inline comments.Jul 9 2022, 3:00 AM
mlir/lib/IR/OperationSupport.cpp
737

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

clementval added inline comments.Jul 11 2022, 4:48 AM
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