Add support for processing Ops with empty region to the CSE

[Ding]

Summary

Automated commit created by applying diff 489654

Phabricator-ID: PHID-HMBT-sx4k6i2amcuosqdga4sn
Review-ID: D141880

[mravishankar]

Thanks! This was the fix I had in mind.

This is where the tests for ops with region are https://github.com/llvm/llvm-project/blob/d2f136920b9247a9e5874d4d3a00a880db6e2827/mlir/test/Transforms/cse.mlir#L327

This is the test op definition https://github.com/llvm/llvm-project/blob/d2f136920b9247a9e5874d4d3a00a880db6e2827/mlir/test/lib/Dialect/Test/TestOps.td#L3141 . You might need to add a new op, or change the op definition to test your case.

mlir/lib/Transforms/CSE.cpp

61–63

A bit mouthful to read... maybe we can just add a helper function

bool hasNoOrEmptyRegion(Operation *op) {
  return op->getNumRegions() == 0 || 
    (op->getNumRegions() == 1 && op->getRegion(0).empty());
}

and use that instead

if (hasNoOrEmptyRegion(lhs) && hasNoOrEmptyRegion(rhs)) {
 ...
}

274–275

I think this would be

if (!hasNoOrEmptyRegion(op)) {
   return failure();
}

?

[Ding]

In D141900#4059778, @mravishankar wrote:

Thanks! This was the fix I had in mind.

This is where the tests for ops with region are https://github.com/llvm/llvm-project/blob/d2f136920b9247a9e5874d4d3a00a880db6e2827/mlir/test/Transforms/cse.mlir#L327

This is the test op definition https://github.com/llvm/llvm-project/blob/d2f136920b9247a9e5874d4d3a00a880db6e2827/mlir/test/lib/Dialect/Test/TestOps.td#L3141 . You might need to add a new op, or change the op definition to test your case.

I found that test.any_cond Op could have empty region, maybe we can use it for testing? Just like this

func.func @cse_ops_with_empty_region()
  -> (i32, i32) {
  %0 = "test.any_cond"() ({
  }) : () -> i32
  %1 = "test.any_cond"() ({
  }) : () -> i32
  return %0, %1 : i32, i32
}

the IR after CSE Pass will be

func.func @cse_single_block_ops_identical_bodies() -> (i32, i32) {
    %0 = "test.any_cond"() ({
    }) : () -> i32
    return %0, %0 : i32, i32
  }

[Ding]

Updating D141900: Add support for processing Ops with empty region to the CSE

[mravishankar]

mlir/lib/Transforms/CSE.cpp

275

WHy cant we use if (!hasNoOrEmptyRegion(op)) here?

mlir/test/Transforms/cse.mlir

484

Nit: Please add newline.

[mravishankar]

Looks good. THanks. Just a few nits.

[rriddle]

mlir/lib/Transforms/CSE.cpp

63

Why just one region? Why not check if all regions are empty?

[Ding]

Updating D141900: Add support for processing Ops with empty region to the CSE

[mravishankar]

Thanks!

[rriddle]

mlir/lib/Transforms/CSE.cpp

44

Please pull these out of the namespace, static functions should be top level.

46

You shouldn't need the parentheses around the all_of call.

[Ding]

Updating D141900: Add support for processing Ops with empty region to the CSE

[mehdi_amini]

Is there anything blocking this from landing?

[mravishankar]

I don't think so. I am AFK. I can merge when I get to my machine

[mravishankar]

@Ding I tried to rebase this and land. Seems to fail tests.

[Ding]

In D141900#4117192, @mravishankar wrote:

@Ding I tried to rebase this and land. Seems to fail tests.

This commit is conflict with the latest commit which support multiple region CSE, I will resolve these conflicts this weekend.

[Ding]

In D141900#4117192, @mravishankar wrote:

@Ding I tried to rebase this and land. Seems to fail tests.

@mravishankar, I found that commit aa4e54f2f444c266310105fccbdbb07fc71894da has fixed this issue when trying to support CSE ops with multiple regions. Should I just push this commit with the new test I added in cse.mlir?

[mravishankar]

Sure. Just push the test. Thanks!

[Ding]

In D141900#4120390, @mravishankar wrote:

Sure. Just push the test. Thanks!

OK, the commit with test is here(https://reviews.llvm.org/D143839).

mlir/lib/Transforms/CSE.cpp

 
 namespace mlir {
 #define GEN_PASS_DEF_CSE
 #include "mlir/Transforms/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 
+static bool hasNoOrEmptyRegions(Operation *op) {
+  return op->getNumRegions() == 0 ||
+         llvm::all_of(op->getRegions(),
+                      [](auto &region) { return region.empty(); });
+}
+
+static bool hasSingleRegionWithSingleBlock(Operation *op) {
+  return op->getNumRegions() == 1 && llvm::hasSingleElement(op->getRegion(0));
+}
+
 namespace {
 struct SimpleOperationInfo : public llvm::DenseMapInfo<Operation *> {
   static unsigned getHashValue(const Operation *opC) {
     return OperationEquivalence::computeHash(
         const_cast<Operation *>(opC),
         /*hashOperands=*/OperationEquivalence::directHashValue,
         /*hashResults=*/OperationEquivalence::ignoreHashValue,
         OperationEquivalence::IgnoreLocations);
   }
   static bool isEqual(const Operation *lhsC, const Operation *rhsC) {
     auto *lhs = const_cast<Operation *>(lhsC);
     auto *rhs = const_cast<Operation *>(rhsC);
     if (lhs == rhs)
       return true;
     if (lhs == getTombstoneKey() || lhs == getEmptyKey() ||
         rhs == getTombstoneKey() || rhs == getEmptyKey())
       return false;
 
-    // If op has no regions, operation equivalence w.r.t operands alone is
-    // enough.
-    if (lhs->getNumRegions() == 0 && rhs->getNumRegions() == 0) {
+    // If op has no or empty regions, operation equivalence
+    // w.r.t operands alone is enough.
+    if (hasNoOrEmptyRegions(lhs) && hasNoOrEmptyRegions(rhs)) {
       return OperationEquivalence::isEquivalentTo(
           const_cast<Operation *>(lhsC), const_cast<Operation *>(rhsC),
           OperationEquivalence::exactValueMatch,
           OperationEquivalence::ignoreValueEquivalence,
           OperationEquivalence::IgnoreLocations);
     }
 
     // If lhs or rhs does not have a single region with a single block, they
     // aren't CSEed for now.
-    if (lhs->getNumRegions() != 1 || rhs->getNumRegions() != 1 ||
-        !llvm::hasSingleElement(lhs->getRegion(0)) ||
-        !llvm::hasSingleElement(rhs->getRegion(0)))
+    if (!hasSingleRegionWithSingleBlock(lhs) ||
+        !hasSingleRegionWithSingleBlock(rhs))
       return false;
 
     // Compare the two blocks.
     Block &lhsBlock = lhs->getRegion(0).front();
     Block &rhsBlock = rhs->getRegion(0).front();
 
     // Don't CSE if number of arguments differ.
     if (lhsBlock.getNumArguments() != rhsBlock.getNumArguments())
LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
   if (isOpTriviallyDead(op)) {
     opsToErase.push_back(op);
     ++numDCE;
     return success();
   }
 
   // Don't simplify operations with nested blocks. We don't currently model
   // equality comparisons correctly among other things. It is also unclear
   // whether we would want to CSE such operations.
-  if (op->getNumRegions() != 0 &&
-      (op->getNumRegions() != 1 || !llvm::hasSingleElement(op->getRegion(0))))
+  if (!hasSingleRegionWithSingleBlock(op) && !hasNoOrEmptyRegions(op))
     return failure();
 
   // Some simple use case of operation with memory side-effect are dealt with
   // here. Operations with no side-effect are done after.
   if (!isMemoryEffectFree(op)) {
     auto memEffects = dyn_cast<MemoryEffectOpInterface>(op);
     // TODO: Only basic use case for operations with MemoryEffects::Read can be
     // eleminated now. More work needs to be done for more complicated patterns

mlir/test/Transforms/cse.mlir

func.func @failing_issue_59135(%arg0: tensor<2x2xi1>, %arg1: f32, %arg2 : tensor<2xi1>) -> (tensor<2xi1>, tensor<2xi1>) {
   %93 = arith.maxsi %false_2, %true_5 : i1
   return %9, %15 : tensor<2xi1>, tensor<2xi1>
 }
 // CHECK-LABEL: func @failing_issue_59135
 //       CHECK:   %[[TRUE:.+]] = arith.constant true
 //       CHECK:   %[[OP:.+]] = test.cse_of_single_block_op
 //       CHECK:     test.region_yield %[[TRUE]]
 //       CHECK:   return %[[OP]], %[[OP]]
+
+// Operations with empty region CSE
+// CHECK-LABEL: func @cse_ops_with_empty_region
+//       CHECK:   %[[OP0:.*[0]]] = "test.any_cond"
+//   CHECK-NOT:   %[[OP1:.*[0]]] = "test.any_cond"
+//        CHECK:   return %[[OP0]], %[[OP0]]
+func.func @cse_ops_with_empty_region()
+  -> (i32, i32) {
+  %0 = "test.any_cond"() ({
+  }) : () -> i32
+  %1 = "test.any_cond"() ({
+  }) : () -> i32
+  return %0, %1 : i32, i32
+}